Why this matters when you build with AI
A dynamic programming table often has a transition that adds up a range of earlier cells. "Count the ways to end here, using any of the last K cells" is a sum over a window. An AI tool writes that sum as a loop over the window, which is correct and easy to read. The cost is K steps for every cell. With n = 10^5 cells and K = 2000, that is 2 * 10^8 steps, which is slow in a judge with a strict time limit, and it grows with K even when the answer does not need K.
The fix is to ask what the window sum is made of. A sum over a window is the difference of two prefix sums, so it costs one subtraction. A window that moves one step at a time can be updated with one addition and one subtraction. These two facts remove the factor K. This lesson shows both, on one counting problem, and checks the results against a brute force program.
A prompt that asks the tool to name the cost before writing code, so you can see whether it knows the window is the issue:
Write C++17 code for this counting problem. The first line has n and K. The second line has n
integers a_i with 0 <= a_i <= 9. Cut the array into consecutive blocks, each of length 1 to K,
so that the sum of every block is divisible by 3. Count the cuts modulo 1000000007.
Before the code, state the cost in terms of n and K, and say whether the inner loop over the
last K values can be replaced by a prefix sum or a running window. Then give the output for the
sample "6 3" and "1 2 3 3 1 2". The output is 5.A good answer names the O(n K) cost, explains how the window sum becomes a difference of two sums or a running total, and prints 5. A bad answer says the code is O(n) and does not explain why, or prints a number without checking it. Check the claim against the brute force in this lesson.
Terms used
- A prefix sum is
Q[i], the sum of the first i values. The sum of values from position lo to position hi - 1 isQ[hi] - Q[lo]. - A window is a range of earlier positions that the current cell reads, such as the last K positions.
- A running window keeps the sum of the current window in a variable. When the window moves one step, the program adds the value that enters and subtracts the value that leaves.
- A difference array stores changes at the ends of ranges. A prefix sum of that array gives the value at each position. It turns "add 1 to a range" into two updates.
- A residue class is the set of positions with the same remainder. Here the remainder is the prefix sum of the array taken mod 3.
- A brute force program checks every possibility. It is slow, but simple enough to trust.
The idea: a window sum is a difference of two prefix sums
Start with a counting problem that has no condition on the blocks. Count the ways to write a total of n as an ordered list of parts, where each part is between 1 and K. Call the count dp[i] for the total i. The last part is some value between 1 and K, so
dp[i] = dp[i - 1] + dp[i - 2] + ... + dp[i - K]
for every i from 1 to n, with dp[0] = 1 and dp[j] = 0 for j below 0. The sum on the right is over a window of K values.
Work out a small case by hand. With K = 2, dp[0] = 1, dp[1] = dp[0] = 1, dp[2] = dp[1] + dp[0] = 2, dp[3] = dp[2] + dp[1] = 3, dp[4] = dp[3] + dp[2] = 5. The answer for n = 4 is 5, which matches the five ways to write 4 as 1 + 1 + 1 + 1, 1 + 1 + 2, 1 + 2 + 1, 2 + 1 + 1 and 2 + 2.
The window sum has the form Q[i] - Q[i - K], where Q[i] is the sum of dp[0] to dp[i - 1]. Every term in the window is in Q[i] and not in Q[i - K], and the terms in Q[i - K] are exactly the ones before the window. So each cell costs one subtraction once Q is known, and Q grows by one addition per cell. The total cost is O(n).
Now add a condition. Each block must have a sum divisible by 3. Let P[i] be the prefix sum of the array modulo 3, with P[0] = 0. A block that covers items j + 1 to i has a sum divisible by 3 exactly when P[j] = P[i]. The condition moved from the block sum to the two endpoints. The transition is
dp[i] = sum over j in [i - K, i - 1] with P[j] = P[i] of dp[j]
Two sums are still needed, but now they must match a class. There are three classes, 0, 1 and 2. Keep one prefix sum per class: Q[i][r] is the sum of dp[j] for j below i with P[j] = r. Then dp[i] = Q[i][P[i]] - Q[lo][P[i]], where lo = max(0, i - K).
Here is the sample, 6 3 and the array 1 2 3 3 1 2. The table shows the prefix residues, the valid earlier positions inside the window and the result.
i P[i] valid j (in the window, same residue) dp[i]
0 0 - 1
1 1 [] 0
2 0 [0] 1
3 0 [0, 2] 2
4 0 [2, 3] 3
5 1 [] 0
6 0 [3, 4] 5Check dp[6] by hand. Position 6 has P = 0. The window for i = 6 with K = 3 is j from 3 to 5. Position 3 has P = 0 and dp = 2, position 4 has P = 0 and dp = 3, and position 5 has P = 1, so it is not counted. The sum is 2 + 3 = 5, which matches the sample answer.
The answer is 0 for many inputs. dp[n] needs P[n] = P[0] = 0, because the blocks cover the whole array, so the total sum must be divisible by 3. For a random array of length 2000 the total sum was not divisible by 3, and every program printed 0. Test on inputs where the total is divisible by 3 as well, otherwise a broken program can pass by printing 0.
Method A: the direct transition
The direct version loops over the window for every cell. It is the code an AI tool usually writes.
// Divisible blocks, method A: direct transition, O(n * K).
#include <bits/stdc++.h>
using namespace std;
const long long MOD = 1000000007LL;
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int n, K;
cin >> n >> K;
vector<int> a(n + 1, 0);
for (int i = 1; i <= n; i++) cin >> a[i];
// P[i] = (a_1 + ... + a_i) mod 3. A block (j, i] has sum divisible by 3 exactly when P[j] == P[i].
vector<int> P(n + 1, 0);
for (int i = 1; i <= n; i++) P[i] = (P[i - 1] + a[i]) % 3;
// dp[i] = number of ways to cut the first i items. dp[0] = 1.
vector<long long> dp(n + 1, 0);
dp[0] = 1;
for (int i = 1; i <= n; i++) {
for (int len = 1; len <= K && len <= i; len++) {
int j = i - len; // the block is items j+1 .. i
if (P[j] == P[i]) dp[i] = (dp[i] + dp[j]) % MOD;
}
}
cout << dp[n] << "\n";
return 0;
}On the sample it prints 5. The inner loop runs min(K, i) times for each i, so the cost is about n * K steps. On a random input with n = 100000 and K = 2000 the loop does about 2 * 10^8 steps. It finished in about 0.4 seconds on my machine. With K = n = 200000 it did not finish within 15 seconds, because the steps are about n^2 / 2 = 2 * 10^10.
Method B: a running window with one bucket per residue
The window moves one step each time i grows by one. The program keeps bucket[r], the sum of dp[j] over the current window with P[j] = r. When i grows, the program does two things. It removes the position that leaves the window, index i - 1 - K, from its bucket. Then it reads the bucket of P[i] as dp[i], and adds dp[i] to its own bucket, so that it is in the window for the next cell.
// Divisible blocks, method B: running window with one bucket per prefix residue, O(n).
#include <bits/stdc++.h>
using namespace std;
const long long MOD = 1000000007LL;
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int n, K;
cin >> n >> K;
vector<int> a(n + 1, 0);
for (int i = 1; i <= n; i++) cin >> a[i];
vector<int> P(n + 1, 0);
for (int i = 1; i <= n; i++) P[i] = (P[i - 1] + a[i]) % 3;
vector<long long> dp(n + 1, 0);
dp[0] = 1;
// bucket[r] = sum of dp[j] over the window j in [i - K, i - 1] with P[j] == r.
long long bucket[3] = {0, 0, 0};
bucket[P[0]] = 1; // j = 0 is in the window for i = 1
for (int i = 1; i <= n; i++) {
int leaving = i - 1 - K; // this index drops out of the window when i grows to i
if (leaving >= 0) {
bucket[P[leaving]] = (bucket[P[leaving]] - dp[leaving] + MOD) % MOD;
}
dp[i] = bucket[P[i]];
bucket[P[i]] = (bucket[P[i]] + dp[i]) % MOD; // dp[i] joins the window for the next i
}
cout << dp[n] << "\n";
return 0;
}On the sample it prints 5. The two updates per cell are the whole cost, so the run time is O(n) for any K. On the random input with n = 100000 and K = 2000 it printed the same answer as method A in about 0.1 seconds, and with K = n = 200000 it finished in the same time.
The order of the two steps matters. The removal must happen before dp[i] is read, because index i - 1 - K is outside the window for i. Section "Common mistakes" shows what happens when the removal uses the wrong index.
Method C: prefix sums per residue class
Method C does the same work with prefix sums instead of a running total. It keeps the full table Q[i][r], which takes one row of three numbers per position. The window sum is then a difference of two entries.
// Divisible blocks, method C: prefix sums of dp per residue class, O(n).
#include <bits/stdc++.h>
using namespace std;
const long long MOD = 1000000007LL;
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int n, K;
cin >> n >> K;
vector<int> a(n + 1, 0);
for (int i = 1; i <= n; i++) cin >> a[i];
vector<int> P(n + 1, 0);
for (int i = 1; i <= n; i++) P[i] = (P[i - 1] + a[i]) % 3;
vector<long long> dp(n + 1, 0);
dp[0] = 1;
// Q[i][r] = sum of dp[j] for j < i with P[j] == r.
// The window sum for i is Q[i][r] - Q[lo][r], where lo = max(0, i - K) and r = P[i].
vector<array<long long, 3>> Q(n + 2, array<long long, 3>{0, 0, 0});
Q[1][P[0]] = dp[0];
for (int i = 1; i <= n; i++) {
int lo = max(0, i - K);
dp[i] = (Q[i][P[i]] - Q[lo][P[i]] + MOD) % MOD;
Q[i + 1] = Q[i];
Q[i + 1][P[i]] = (Q[i][P[i]] + dp[i]) % MOD;
}
cout << dp[n] << "\n";
return 0;
}The subtraction can go negative after the modulo, so it adds MOD before the second modulo. Forgetting that step gives a negative answer on some inputs, and the printed value has a minus sign. Method B needs the same care.
Difference arrays: add to a range in O(1)
Prefix sums also help in the other direction. Suppose the DP wants to add a value to every position from l to r, and read the values at the end. Doing it one position at a time costs the length of the range for each update. A difference array stores +1 at position l and -1 at position r + 1. Then one prefix sum pass gives every position its total.
Here is an example that counts, for each position, how many of m ranges cover it. The ranges are 1 to 3, 2 to 6 and 6 to 6, on an array of length 6.
// Count, for each position, how many of the m ranges [l, r] cover it.
#include <bits/stdc++.h>
using namespace std;
int main() {
int n, m;
cin >> n >> m;
vector<long long> diff(n + 2, 0);
for (int k = 0; k < m; k++) {
int l, r;
cin >> l >> r;
diff[l] += 1; // +1 starts at l
diff[r + 1] -= 1; // and stops after r
}
long long run = 0;
for (int i = 1; i <= n; i++) {
run += diff[i]; // the prefix sum of the difference array gives the cover count
cout << run << (i == n ? "\n" : " ");
}
return 0;
}The sample input 6 3, then 1 3, 2 6 and 6 6, prints 1 2 2 1 1 2. A brute force that marks each position in each range prints the same line. I also compared the two on 200 random sets of ranges, and they matched on all of them. A DP uses the same idea when a transition adds one value to a whole range of later cells.
Complexity and when each method is worth it
Method A costs O(n K) time and O(n) memory. Method B costs O(n) time and O(n) memory, and it only needs the last K entries of dp at any time, so the memory can be cut to O(K) by a ring buffer if the values are not needed later. Method C costs O(n) time and O(n) memory, and it is the easiest to check by eye, because each window sum is a visible subtraction. For a contest, method B is the usual choice, and method C is good when you want a version that is easy to explain.
Use method A on small inputs as the reference in your stress test. It is short, and a bug in the window code shows up as a difference from it.
Common mistakes
Off-by-one in the removal index. If method B removes index i - K instead of i - 1 - K, the window becomes one value too short. On the sample it prints 2 where the correct answer is 5. The real output from my run:
blocks_window (correct): 5
blocks_window with leaving = i - K: 2On 200 random small inputs the buggy version gave a wrong answer on 27 of them. The window for cell i must include index i - K, because the block from i - K + 1 to i has length K.
Reading the bucket before removing the leaving index. The order is part of the transition. Read the bucket after the removal, as the code does.
Negative values after the modulo. The subtraction bucket - dp + MOD must be taken before the second modulo. A plain (bucket - dp) % MOD gives a negative number in C++ when the difference is below zero, and the next addition carries the minus sign into the answer.
Using the wrong class for the window sum. In method C the difference must use the class of P[i], both for the old and new prefix sums. Subtracting a different class gives a value that has nothing to do with the transition.
Testing only on answers that are 0. A random array often gives 0 because the total sum is not divisible by 3. A wrong program can print 0 by accident. Use inputs where the total sum is divisible by 3, and compare against the brute force on those.
Memory for large n. Method C stores n + 2 rows of three 64-bit numbers, about 4.8 MB for n = 200000. That is fine for most limits, but method B needs far less.
Checklist
- Write the cost of the direct transition, in terms of n and the window length, before choosing a method.
- Check whether the window sum is a range of one array. If it is, use a prefix sum and one subtraction.
- If the window moves one step at a time, keep a running total and update it with one addition and one subtraction.
- For a condition on residues, keep one total per class and use the class of the current position.
- Add MOD before the second modulo in every subtraction.
- Test against the direct transition on small random inputs, and include inputs where the answer is not 0.
- Check the removal index on a window of length 1 and a window of length K.
Summary
Most window transitions are a sum over a range of earlier cells, and a range sum needs one subtraction once the prefix sums exist. The running window and the prefix-sum table give the same answer at O(n) cost, and the difference array runs the same idea in the other direction. Write the direct version first, then replace the inner loop, and compare the two on small inputs.
Practice problems
- Warm-up. Count the ways to write a total of n as an ordered sum of parts from 1 to K, with
dp[0] = 1anddp[i]the sum of the last K values. Write the formula for dp[i] using a prefix sumQ, and give the answer for n = 10 and K = 3.
Show the answer
Answer: with Q[i] the sum of dp[0] to dp[i - 1], dp[i] = Q[i] - Q[max(0, i - K)]. The values for K = 3 are dp = 1, 1, 2, 4, 7, 13, 24, 44, 81, 149, 274, so the answer for n = 10 is 274. I checked this with a short Python loop. The same loop with the window sum written out gives 5 for n = 4 and K = 2.
- Coverage. Use the difference array program on the input
5 2with ranges1 5and3 3. Predict the output before you run it.
Show the answer
Answer: 1 1 2 1 1. Position 3 is covered by both ranges, so its count is 2. The other positions are covered by the first range only. Run the program to confirm the line.
- Find the off-by-one. Explain why
leaving = i - Kgives 2 on the sample, not 5. Use the window for i = 6 with K = 3.
Show the answer
Answer: with leaving = i - K, the index 3 is removed when i = 6, before dp[6] is read. The correct window for i = 6 is j from 3 to 5, so index 3 must stay. The window used by the buggy version is only j = 4 and 5, which is one value too short. The sum loses the term from j = 3 (dp = 2 with P = 0), and the answer changes.
- Change the condition. Change the rule so that every block must have an even sum, and keep K = 3 and the sample array. Which lines change in method B, and what is the answer?
Show the answer
Answer: the modulus 3 becomes 2 in the prefix sums and in the bucket array, and the bucket array now has two entries instead of three. The answer on the sample is 2, which I checked with a Python program that counts every cut. The window logic stays the same. This is the same pattern with a smaller number of classes.