State design and transitions: count the work before you optimise

Why this matters when you build with AI

Ask an AI tool for a dynamic programming (DP: a table of answers to smaller problems, where the final cell is the answer) solution to a partition problem and you will often get code that is right on the sample and too slow on the judge. The usual cause is that nobody counted the work. The tool writes three nested loops, the sample has six numbers, and the program finishes in a few milliseconds. At n = 2000 the same loops run more than a billion steps.

The fix is to count before you run anything. Write down what one table cell means, how many cells there are, and how many earlier cells each one reads. Multiply the last two numbers. Only then ask whether the count fits your limit. This lesson practises that order on one problem, then measures the result on a real machine.

A prompt that asks for the count in the answer, so a wrong state or a missing count shows up early:

Write C++17 code for this problem. The first line has n and K. The second line has n integers.
Split the integers into exactly K non-empty contiguous groups. A group with L items and item sum S
costs L * L + S % 10. Minimise the total cost. Use long long.
Before the code, state the meaning of the DP table in one sentence, give the transition,
the order in which cells are filled, and the number of states times the cost of one transition.
Then give the code and its output on the sample "6 3" and "1 2 3 4 5 6". The output starts with 23.

A good answer states what each table cell means, gives a formula for the count, and prints 23 on the sample. A bad answer has table indexes with no stated meaning, or a complexity claim that the code does not match. Check it against the brute force in the worked example before you trust it.

Terms used

• A partition here means splitting the row of n items into groups, so every item is in exactly one group.
• A group is contiguous when its items are next to each other in the row. Items 3, 4 and 5 form a contiguous group. Items 3, 5 and 7 do not.
• A state is one cell of the DP table. Its meaning is a question the cell answers, for example "what is the cheapest way to split the first 4 items into 2 groups?"
• A transition is the rule that computes one state from earlier states.
• A layer is all states with the same group count g. Layer g is built from layer g - 1.
• A rolling array keeps only the last layer instead of the whole table.
• A choice table stores, for each state, which earlier state it came from. You need it to print the groups.
• A brute force tries every possible answer. It is too slow for large inputs, but it is simple enough to trust.

The idea: a state is a question you can answer

The problem: you get n integers a_1 to a_n and a number K. Split them into exactly K non-empty contiguous groups. A group with L items and item sum S costs L * L + S % 10, where % 10 means the last digit of S. Find the smallest total cost of all groups.

Start with an example. Take the integers 1 2 3 4 5 6 and K = 3. One valid split is [1 2], [3 4], [5 6]. The first group has L = 2 and S = 3, so its cost is 4 + 3 = 7. The second has S = 7, cost 4 + 7 = 11. The third has S = 11, cost 4 + 1 = 5. The total is 23. The DP has to find the cheapest split, and it must do that without trying every split.

The state is the question the table answers. Let dp[g][i] be the smallest total cost for splitting the first i items into exactly g groups. Fix this meaning first. Every transition below is checked against it.

The transition asks where the last group starts. Say the last group covers items j + 1 to i. The first j items must form g - 1 groups, and the cheapest way to do that is dp[g-1][j], so

dp[g][i] = the minimum over j of dp[g-1][j] + cost(items j+1 to i)

The base case is dp[0][0] = 0, with every other dp[0][i] unreachable. The answer is dp[K][n].

Here is the table for the sample. Each column is a value of i, from 1 to 6. A dash means that the cell cannot be reached, for example 1 item cannot form 2 groups.

i      1   2   3   4   5   6
g=1    2   7   15  16  30  37
g=2    -   5   11  18  18  21
g=3    -   -   9   16  16  23

Check two cells by hand. The g = 1 cell for i = 6 is one group with L = 6 and S = 21, so 36 + 1 = 37. The g = 2 cell for i = 3 is the best of two choices. Starting the last group at item 2 gives dp[1][1] + cost(items 2 to 3) = 2 + (4 + 5 % 10) = 2 + 9 = 11. Starting it at item 3 gives dp[1][2] + cost(item 3) = 7 + (1 + 3) = 11. Both are 11, which matches the table. The answer dp[3][6] = 23 matches the sample.

Order the transitions

Each cell of layer g reads only layer g - 1, and only cells with a smaller i. So the order is: loop over g from 1 to K on the outside, and over i on the inside. Finish all of layer g - 1 before you start layer g.

If you swap the loops, a cell reads values that are not computed yet. The code does not crash. It reads the starting value instead of a computed one, so the answer can be wrong. Writing the order down as a sentence before the code is the cheapest check against this.

Count states times transition cost before optimising

The table has about K * n states, since g runs from 1 to K and i from g to n. Each state tries every j between g - 1 and i - 1. The exact number of inner steps is the sum over g and i of (i - g + 1). This is a count, not an estimate, and it depends only on n and K.

The count is the number to compare with your limit. A table with 10^6 states that each try 10^3 choices is 10^9 steps, which is usually too many. A table with 10^4 states that each try 10^2 choices is 10^6 steps, which is almost always fine.

Here is the count for the naive code on random inputs with K = n / 2, measured on the machine described below:

The step counts are exact, from a formula that matches the loops in the code. The times are from one machine: 12 hardware threads, 20 GB of memory, Linux under WSL2, g++ 15.2.0 with -O2 -std=c++17. Your times will differ, but the step counts will not. Compare the step count with your limit, and treat the time as a rough guide only.

The table shows where the naive version breaks. Going from n = 1200 to n = 2000 multiplies the steps by about 4.6, and the time by about 3.4. Later lessons in this module remove the inner loop. This lesson does not optimise anything. It shows the table, the order and the count, so that you know what you are optimising later.

Dimension reduction and rolling arrays

The full table has (K + 1) * (n + 1) cells. Layer g reads only layer g - 1, so you need two rows, not K + 1. Keep them as prev and cur, and swap them after each layer:

vector<ll> prev(n + 1, INF), cur(n + 1, INF);
prev[0] = 0;
// ... fill cur from prev for each g ...
swap(prev, cur);

This reduces the value memory from (K + 1) * (n + 1) long long values to 2 * (n + 1). The time does not change, because the same transitions still run.

Rolling rows lose the history, and printing the groups needs the history. The fix is a second table, choice[g][i], that stores the chosen j for each state. It has (K + 1) * (n + 1) ints. For n = 2000 and K = 1000 that is about 2.0 million ints, or 8 MB. The walk back from the answer is: start at i = n, and for g from K down to 1 read j = choice[g][i], output the group j + 1 to i, then set i = j.

Memory is a limit too. If you keep the full table of long long values for n = 2000 and K = 1000, it takes 1001 * 2001 * 8 bytes, which is about 16 MB. Check the memory limit of your judge in the same way you check the time limit.

Worked example: the full program

This is the naive program. It runs the transitions in the order above, rolls the value rows, keeps a choice table, and prints the cost and the groups. The file compiles with g++ -O2 -std=c++17.

// Lesson 01: split n items into exactly K contiguous groups, minimise total box cost.
// Box cost of a group with L items and sum S is L*L + S % 10.
// Plain DP: O(K * n^2) time, one rolling row per layer, choice table for the groups.
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const ll INF = LLONG_MAX / 4;
 
static inline ll boxCost(ll len, ll sum) { return len * len + sum % 10; }
 
int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int n, K;
    cin >> n >> K;
    vector<ll> a(n);
    for (auto &x : a) cin >> x;
    if (K > n) {
        cout << -1 << "\n";
        return 0;
    }
 
    vector<ll> pre(n + 1, 0);
    for (int i = 0; i < n; i++) pre[i + 1] = pre[i] + a[i];
 
    // prev[i] = best cost of splitting the first i items into g-1 groups.
    // cur[i]  = the same for g groups. Only two rows are kept.
    vector<ll> prev(n + 1, INF), cur(n + 1, INF);
    prev[0] = 0;
    // choice[g][i] = end of the first g-1 groups, so the g-th group is choice+1..i.
    vector<vector<int>> choice(K + 1, vector<int>(n + 1, -1));
 
    for (int g = 1; g <= K; g++) {
        fill(cur.begin(), cur.end(), INF);
        for (int i = g; i <= n; i++) {
            for (int j = g - 1; j < i; j++) {
                if (prev[j] == INF) continue;
                ll c = prev[j] + boxCost(i - j, pre[i] - pre[j]);
                if (c < cur[i]) {
                    cur[i] = c;
                    choice[g][i] = j;
                }
            }
        }
        swap(prev, cur);
    }
 
    cout << prev[n] << "\n";
 
    // Walk the choices back from n to recover the K groups (1-indexed, inclusive).
    vector<pair<int, int>> groups;
    int i = n;
    for (int g = K; g >= 1; g--) {
        int j = choice[g][i];
        groups.push_back({j + 1, i});
        i = j;
    }
    reverse(groups.begin(), groups.end());
    for (auto [l, r] : groups) cout << l << " " << r << "\n";
    return 0;
}

Output on the sample 6 3 and 1 2 3 4 5 6:

23
1 2
3 4
5 6

The sum 7 + 11 + 5 = 23 matches the hand calculation at the start of the lesson.

Check it against a brute force

The brute force tries every set of cut points. There are 2^(n-1) sets of cuts, so it is only for small n. It keeps the sets that give exactly K groups and returns the smallest cost. The same check in Python, with the sample hard-coded, is below. It is the same idea as the C++ brute force, using combinations to pick K - 1 cut points.

Output:

23
at most 3 groups, by mistake: 21
K=1..6: [37, 21, 23, 21, 19, 27]

The C++ program and the brute force were compared on 1500 random inputs with n from 1 to 11, values from 1 to 30, and K from 1 to n + 2. The two agreed on every input. For each DP answer the test also checked that the printed groups are contiguous, cover all items, number exactly K, and add up to the printed cost. For the cases with K greater than n, both print -1. This test was run on the machine described above, and it shows that the code is right for these sizes. It does not prove that it is right for every input.

Measure where it breaks

Build the input with random values up to 10^6 and K = n / 2, then count the steps and time the run. The step formula is the same as the loops:

def steps(n, K):
    return sum(i - g + 1 for g in range(1, K + 1) for i in range(g, n + 1))

On a judge the number to compare is the step count and the limit for your language and machine. Pick a limit you have measured for your own machine, not one you were told. Note the machine (CPU, compiler, flags) next to the number, so that you can compare runs later.

Common mistakes

1. Counting "at most K" groups instead of "exactly K". If you take the minimum over every layer, the sample 6 3 / 1 2 3 4 5 6 prints 21 (the cost of 2 groups), not 23. Read only dp[K][n].

2. Forgetting to reset the current row. If the fill at the start of each layer is removed, a cell can keep a value from two layers back. For 6 6 with the values 12 26 23 27 24 21 the code prints 21. The right answer is 29, the cost of six single items: 3 + 7 + 4 + 8 + 5 + 2. In 104 of 400 random inputs with n from 2 to 9 this version disagreed with the brute force.

3. Overflow. A sum of 5000 values, each up to 10^6, can reach 5 * 10^9. That is larger than 2147483647, the largest value of a 32-bit int. Use long long for the prefix sums and the DP values, as the code does.

4. Reconstructing from the rolling rows. The two rows hold only the last layer, so you cannot walk back through them. If you try to rebuild the groups from the rolled rows, the cost can be right while the groups are not.

5. Swapping the loop order. A layer must be complete before the next one starts, because each cell reads the previous layer. See the section on order.

Checklist

• Write the meaning of the DP state in one sentence before any code.
• Write the transition and the base case, and check the base case against the sample.
• Fix the order of the layers, and check that each cell reads only finished cells.
• Count states times transitions, and compare the count with your limit.
• Use long long for sums and DP values, and check the largest possible sum.
• Keep a choice table if you have to print the answer, not only its value.
• Test against a brute force on small random inputs, and check the printed answer by its own rules, not only by its cost.

Summary

A DP is a table with a meaning for each cell, a rule for each transition, and an order for the cells. Before you optimise, count the cells and the transitions and compare the total with your limit. In this lesson the count is the exact sum given above, and it is measured in the table. The rolling row saves memory without changing the time, and the choice table keeps the answer you need to print.

Practice

1. Count the steps. How many inner-loop steps does the naive program run for n = 10 and K = 3?

Show the answer

Answer: 136. The inner loop runs j from g - 1 to i - 1, which is i - g + 1 steps. Sum that over g from 1 to 3 and i from g to 10, using the formula in the section "Measure where it breaks".

Hint: write the sum for each g on its own line, then add the three totals.

2. Change the cost. Replace the box cost with L * L, with no S % 10 term. Which lines of the program change, and what is the answer on the sample?

Show the answer

Answer: only boxCost changes. The loops, the rolling rows and the choice table stay the same. On the sample the answer is 12, from three groups of two items each, which is 4 + 4 + 4. Using equal group sizes is the best here because the sum of squares is smallest when the sizes are as equal as possible.

Hint: find every place where the cost is used. There is one call in the loop, and the program keeps all of its other lines.

3. At most K. The task now says "use between 1 and K groups". Write the final step that gives the answer, and say what it prints on the sample.

Show the answer

Answer: take the minimum over g from 1 to K of dp[g][n], and print it. On the sample the values are 37, 21 and 23, so the minimum is 21, which is the output of the brute force line "at most 3 groups". The rest of the program does not change.

Hint: the answer for each g is in the last column of that layer. Keep the minimum while you fill the layers, or keep the last column of each layer in a list.

4. Memory. How many bytes does the full table of long long values take for n = 2000 and K = 1000, with (K + 1) * (n + 1) cells of 8 bytes each?

Show the answer

Answer: 16,024,008 bytes, about 16 MB. The count is 1001 * 2001 = 2,003,001 cells, and 2,003,001 * 8 = 16,024,008.

Hint: multiply the two dimensions first, then the size of one value in bytes.

5. Keep the history. A teammate wants to replace the choice table with one row, the way the value rows are rolled. Explain what breaks, and give a case where the walk back gives the wrong groups.

Show the answer

Answer: one row holds only the choice for the last layer that was filled. The walk back needs the choice for g = K, then g = K - 1, and so on, each at a different i. Once the rows are overwritten, the choice for an earlier layer is gone, so the walk cannot recover those groups. The cost is still right, because the value rows are kept separately, but the groups cannot be printed from one row. Test this by printing the groups and checking that they add up to the cost.

Hint: write the walk back as a list of (g, i) pairs. Check whether any two pairs share a g.