Why this matters when you build with AI
A counting problem often asks "how many ways give total s?" for every s. An AI tool that is asked for this usually writes a dynamic programming loop over items and totals, which is correct and takes time proportional to items times total. When the total reaches 200000 and there are 2000 items, that loop does about 4 times 10^8 steps, and the judge says "time limit exceeded" with no hint of why. The fix, polynomial multiplication with the NTT, is a different algorithm, and a tool that does not know it will not suggest it.
AI tools also make two specific mistakes with this topic. The first is the modulus. Many snippets on the internet use 10^9 + 7, and that prime cannot be used with the NTT in the form below, because the transform needs a root of unity of a length that 10^9 + 7 does not have. The second is the sign of the output. A tool may print the transform of the product without the final division by the length, and the numbers look plausible but are off by a fixed factor. A small brute force over random inputs catches both problems in seconds.
This lesson covers ordinary generating functions, the convolution that products of them give, and an iterative NTT in C++17. Every program here was compiled with g++ -std=c++17 -O2 -Wall -Wextra using GCC 15.2.0, and the output blocks are from our runs. The stress test compares the NTT with a quadratic product, and the subset counts with a knapsack loop, on random inputs. Timing numbers depend on your machine.
The idea: counting with polynomials
An ordinary generating function (OGF) for a sequence a_0, a_1, a_2, ... is the polynomial, or power series, with those numbers as coefficients:
A(x) = a_0 + a_1 x + a_2 x^2 + a_3 x^3 + ...
Here x is only a placeholder. We never substitute a number for x. The exponent of x records the total, and the coefficient records how many ways reach that total. This is the whole trick: put each choice in a polynomial, multiply the polynomials, and read the answer off the coefficients.
A concrete example first. Take one die with faces 1 to 6. Each face is one way to get a total, so the generating function for one die is
D(x) = x + x^2 + x^3 + x^4 + x^5 + x^6
The coefficient of x^3 is 1, because only the face 3 gives total 3. Now roll two dice and ask how many ways give total 7. Write the second die's polynomial too, and multiply:
D(x) times D(x) = (x + x^2 + ... + x^6)(x + x^2 + ... + x^6)
When you multiply, each term of the first polynomial meets each term of the second, and the exponents add. The term x^a times x^b gives x^(a+b). So the coefficient of x^7 counts pairs (a, b) with a + b = 7, both between 1 and 6. There are 6 such pairs, so the coefficient is 6. The product rule does the counting for you.
The same idea handles subsets. Take an item with price p. Either you skip it, which contributes 1, or you take it, which contributes x^p. So each item is the factor (1 + x^p). The product over all items is the generating function of subset totals: the coefficient of x^t is the number of subsets whose prices add up to t.
Try four tags with prices 1, 1, 2 and 3. The product (1 + x)(1 + x)(1 + x^2)(1 + x^3) is computed by a short program below.
Output from our run:
x^0: 1
x^1: 2
x^2: 2
x^3: 3
x^4: 3
x^5: 2
x^6: 2
x^7: 1The coefficient of x^3 is 3. The three subsets with total 3 are {3}, {1a, 2} and {1b, 2}, where 1a and 1b are the two tags with price 1. The coefficients add to 16, which is 2 to the power 4, the number of subsets of four tags. The tags are told apart by position, so the two 1-tags count as different.
Products are convolutions
Multiplying two polynomials is a convolution. If A has coefficients a_i and B has coefficients b_j, then the product C = A times B has
c_k = sum over i of a_i times b_(k - i)
This sum runs over every pair with i + j = k. It is the same sum you would write for the number of pairs, and it is the reason the die example worked. The direct way to compute all c_k takes about n times m multiplications for polynomials of lengths n and m. For lengths 10^5 this is 10^10, far too slow.
The generating function view gives a second route. Counting by total is a sequence of steps, one per item, and each step is a convolution with a short polynomial. A DP over items does exactly this with a short polynomial (1 + x^p), and its cost is items times total. Multiplying the full polynomials with a fast convolution can do better, when the factors are many and the total is large.
Two more facts matter for the rest of the lesson.
- Coefficients are counts, so they can be huge. We keep them modulo a prime p, and all arithmetic in the transform is done modulo p.
- Products of many polynomials can be done in any grouping. Multiplying them in pairs, arranged as a balanced tree, costs about S log S per level of the tree, where S is the final length. The number of levels is the logarithm of the number of factors. That is the divide and conquer used in the worked problem.
The NTT
The number theoretic transform (NTT) is the fast Fourier transform (FFT) done with integers modulo a prime, instead of with complex numbers. It turns a polynomial of length n into its values at n points, the n-th roots of unity, using about n log n operations. Multiplying two polynomials then becomes three steps.
- Transform both polynomials to point values.
- Multiply the point values one by one, which costs n operations.
- Transform back, and divide every value by n.
The result is the product, computed in about n log n time instead of n squared. Exactness matters for counting, so we use integers. The idea needs a prime p and a number g such that g has order exactly n modulo p, meaning n is the smallest positive power with g^n = 1 mod p. Such a g is a root of unity of order n.
The prime used in practice is p = 998244353, with g = 3. Here is why it works, checked with a short script.
- p - 1 = 998244352 = 2^23 times 7 times 17.
- So p - 1 is divisible by 2^23 = 8388608, and the transform works for any length that is a power of two up to 2^23.
- The number 3 is a primitive root modulo p: 3 raised to (p - 1) divided by 2, by 7 and by 17 is not 1, so 3 has the largest possible order, p - 1.
- The root of order n is 3 raised to (p - 1) divided by n, when n divides p - 1.
Libraries that do exact convolution use the same prime and root. The AtCoder Library's convolution uses 998244353, and its primitive root function returns 3 for it, as shown in ports of that library such as ac-library.cr convolution.cr and ac-library-rb. We did not read the official AtCoder Library page itself for this lesson, so check it if you need the exact function.
Why not 10^9 + 7? Its p - 1 is 1000000006, which equals 2 times 500000003, and 500000003 is prime. The only power of two that divides p - 1 is 2 itself, so the transform works only for length 2. A product of two long polynomials then has no matching root. A program written for 10^9 + 7 gives wrong numbers, often without any error.
Iterative implementation
The recursive FFT splits the input into even and odd positions. The iterative version does the same work in place. First it reorders the array into bit-reversed order, so that the butterfly steps can run on neighbours. Then for each block size len = 2, 4, 8, ..., it combines pairs of halves. Each pair (u, v) becomes (u + w times v, u - w times v), where w is the root of order len raised to the power k for position k in the half.
The inverse transform uses the inverse root, and at the end every value is multiplied by the inverse of n. Forgetting that last division gives an answer that is n times too large, which is one of the mistakes below.
The code
The file below holds the NTT, a naive product for testing, the subset counting with a divide and conquer product, the knapsack loop for testing, and a stress test. Run it with stress as the argument to test. Without an argument it reads the input from standard input.
#include <bits/stdc++.h>
using namespace std;
using u32 = unsigned int;
using u64 = unsigned long long;
const u32 MOD = 998244353;
u64 pw(u64 b, u64 e) {
u64 r = 1;
b %= MOD;
while (e > 0) {
if (e & 1) r = r * b % MOD;
b = b * b % MOD;
e >>= 1;
}
return r;
}
// In-place iterative NTT. Length must be a power of two dividing 2^23.
void ntt(vector<u32>& a, bool invert) {
int n = a.size();
// Bit-reversal permutation.
for (int i = 1, j = 0; i < n; i++) {
int bit = n >> 1;
for (; j & bit; bit >>= 1) j ^= bit;
j ^= bit;
if (i < j) swap(a[i], a[j]);
}
for (int len = 2; len <= n; len <<= 1) {
u64 wlen = pw(3, (MOD - 1) / len);
if (invert) wlen = pw(wlen, MOD - 2);
int half = len / 2;
vector<u32> w(half);
w[0] = 1;
for (int k = 1; k < half; k++) w[k] = (u64)w[k - 1] * wlen % MOD;
for (int i = 0; i < n; i += len) {
for (int k = 0; k < half; k++) {
u64 u = a[i + k];
u64 v = (u64)a[i + k + half] * w[k] % MOD;
a[i + k] = (u + v) % MOD;
a[i + k + half] = (u + MOD - v) % MOD;
}
}
}
if (invert) {
u64 inv_n = pw(n, MOD - 2);
for (auto& x : a) x = (u64)x * inv_n % MOD;
}
}
// Polynomial product modulo MOD using the NTT.
vector<u32> multiply(vector<u32> a, vector<u32> b) {
int need = a.size() + b.size() - 1;
int sz = 1;
while (sz < need) sz <<= 1;
a.resize(sz);
b.resize(sz);
ntt(a, false);
ntt(b, false);
for (int i = 0; i < sz; i++) a[i] = (u64)a[i] * b[i] % MOD;
ntt(a, true);
a.resize(need);
return a;
}
// Quadratic reference product, used only for tests.
vector<u32> multiply_naive(const vector<u32>& a, const vector<u32>& b) {
vector<u64> c(a.size() + b.size() - 1, 0);
for (size_t i = 0; i < a.size(); i++)
for (size_t j = 0; j < b.size(); j++)
c[i + j] = (c[i + j] + (u64)a[i] * b[j]) % MOD;
vector<u32> r(c.begin(), c.end());
return r;
}
// Product of polynomials ps[lo..hi) by divide and conquer.
vector<u32> product_range(const vector<vector<u32>>& ps, int lo, int hi) {
if (hi - lo == 1) return ps[lo];
int mid = (lo + hi) / 2;
return multiply(product_range(ps, lo, mid), product_range(ps, mid, hi));
}
// Counts of subsets by total price, for prices in [1, 100], modulo MOD.
// Result has size P + 1, where P is the sum of all prices.
vector<u32> count_by_total(const vector<int>& prices) {
const int MAXN = 2005;
vector<u64> fact(MAXN), inv_fact(MAXN);
fact[0] = 1;
for (int i = 1; i < MAXN; i++) fact[i] = fact[i - 1] * i % MOD;
inv_fact[MAXN - 1] = pw(fact[MAXN - 1], MOD - 2);
for (int i = MAXN - 1; i > 0; i--) inv_fact[i - 1] = inv_fact[i] * i % MOD;
auto C = [&](int n, int k) -> u64 {
if (k < 0 || k > n) return 0;
return fact[n] * inv_fact[k] % MOD * inv_fact[n - k] % MOD;
};
vector<int> cnt(101, 0);
for (int p : prices) cnt[p]++;
vector<vector<u32>> ps;
for (int v = 1; v <= 100; v++) {
if (cnt[v] == 0) continue;
int m = cnt[v];
// Choosing j tags of price v adds v * j to the total, with C(m, j) ways.
vector<u32> poly(v * m + 1, 0);
for (int j = 0; j <= m; j++) poly[v * j] = C(m, j);
ps.push_back(poly);
}
return product_range(ps, 0, ps.size());
}
// Quadratic reference: 0/1 knapsack counting over the prices, one item at a time.
vector<u32> count_by_total_naive(const vector<int>& prices) {
int P = accumulate(prices.begin(), prices.end(), 0);
vector<u64> dp(P + 1, 0);
dp[0] = 1;
int cur = 0;
for (int x : prices) {
cur += x;
for (int t = cur; t >= x; t--) dp[t] = (dp[t] + dp[t - x]) % MOD;
}
vector<u32> r(dp.begin(), dp.end());
return r;
}
void stress() {
mt19937 rng(12345);
int cases = 0;
// NTT multiply against the quadratic product.
for (int it = 0; it < 300; it++) {
int la = rng() % 200 + 1, lb = rng() % 200 + 1;
vector<u32> a(la), b(lb);
for (auto& x : a) x = rng() % MOD;
for (auto& x : b) x = rng() % MOD;
if (multiply(a, b) != multiply_naive(a, b)) {
printf("NTT MISMATCH at case %d\n", it);
exit(1);
}
cases++;
}
// Subset counts by total against the knapsack loop.
for (int it = 0; it < 2000; it++) {
int n = rng() % 30 + 1;
vector<int> prices(n);
for (auto& x : prices) x = rng() % 20 + 1;
if (count_by_total(prices) != count_by_total_naive(prices)) {
printf("SUBSET MISMATCH at case %d\n", it);
exit(1);
}
cases++;
}
printf("stress: %d cases, 0 mismatches\n", cases);
}
int main(int argc, char** argv) {
if (argc > 1 && string(argv[1]) == "stress") {
stress();
return 0;
}
int n;
if (!(cin >> n)) return 0;
vector<int> prices(n);
for (auto& x : prices) cin >> x;
vector<u32> res = count_by_total(prices);
for (size_t t = 0; t < res.size(); t++) {
if (t) cout << ' ';
cout << res[t];
}
cout << '\n';
return 0;
}Output from our run, for the four tags above and for the stress test:
$ printf '4\n1 1 2 3\n' | ./subset_by_value
1 2 2 3 3 2 2 1
$ ./subset_by_value stress
stress: 2300 cases, 0 mismatchesThe first line matches the Python expansion above, term for term. The stress test has 300 NTT products of random polynomials with lengths up to 200, and 2000 random price lists of up to 30 prices, each up to 20. Both comparisons agreed on every case.
Worked problem: subset counts by total
Problem. Lena has n price tags. Tag i shows price p_i, with 1 at most p_i at most 100. She picks any subset of the tags, possibly empty. For each total T from 0 to P, where P is the sum of all prices, count the subsets whose prices add up to exactly T. Print the counts modulo 998244353, separated by spaces. The input has n on the first line, then the n prices. Limits: n at most 2000, so P is at most 200000.
Idea. Group the tags by price. If m tags have price v, then choosing j of them adds v times j to the total, in C(m, j) ways. So the generating function for price v is the sum over j of C(m, j) x^(v j). There are at most 100 distinct prices, so there are at most 100 factors. Multiply them with the divide and conquer product, and the coefficient of x^T is the answer.
This is faster than the knapsack loop when many tags share a price. The NTT product of 100 factors takes about S log S per level of the tree, and the total degree S is at most 200000. The knapsack loop takes n times P, which is 4 times 10^8 for the largest input.
Full solution. It is the count_by_total function in the code above. The steps are these.
- Count the tags of each price in
cnt. - For each price v with m tags, build the polynomial with C(m, j) at exponent v j, using factorials and inverse factorials modulo the prime. Since the prime is larger than 2000, no factorial is zero modulo it.
- Multiply all these polynomials with
product_range. The degree of the product is the sum of v times m over all v, which is P, so the result has P + 1 coefficients. - Print the coefficients.
Test. The stress test checks this function on 2000 random price lists against the knapsack loop, which is a separate and simpler method. There were 0 mismatches. The worked sample, four tags with prices 1, 1, 2, 3, gives the same counts as the Python expansion. We also timed the largest input we could think of: 2000 prices, each random between 1 and 100, with P near 10^5. It ran in about 0.1 seconds on our machine. Your times will differ, and the exact total depends on the random seed.
Output for the largest case. The program printed 97396 numbers on that run, which is P + 1 for that seed. The count is not a fixed value, so use it only as a check on the length.
Complexity
Writing the product of d factors of total degree S by divide and conquer takes O(S log S log d) time, since each of the log d levels of the tree multiplies polynomials of total length about S. Memory is O(S). For the worked problem with d at most 100 and S at most 200000, this is a few tens of millions of basic steps.
The NTT itself uses an array of length at least the product length, rounded up to a power of two. For product length up to 2^23 it works with this modulus. The subset product needs about 200001 coefficients, so the array is 2^18 = 262144 long, well inside the limit.
Common mistakes
Using 10^9 + 7 with the NTT. Symptom: the answers are wrong for long inputs, with no error. The code computes (MOD - 1) / len with integer division, and for 10^9 + 7 that division does not give a root of the right order. Use 998244353, or another prime with a suitable power of two in p - 1. We did not run this variant, so treat this as a description of the arithmetic, not a measured result.
Forgetting the division by n. Symptom: every answer is n times too large. The inverse transform must multiply each value by the inverse of the length. The code does this with inv_n.
Too short an array. The product of lengths a and b has length a + b - 1. If you pad to the length of one input, the high coefficients wrap around and add into the low ones. Symptom: small totals come out too large. Pad to at least a + b - 1, rounded up to a power of two.
Overflow in a product. Two values below 998244353 multiply to about 10^18, which fits in 64 bits, but a sum of several such products does not. Reduce after each multiplication, as the code does. Using int for the values is a common source of silent overflow.
Mixing up ordinary and exponential generating functions. The ordinary form used here puts a_k in the coefficient of x^k. The exponential form divides by k!. A problem about orderings or labelled objects often needs the exponential form, so check the statement before you copy a formula.
Trusting one small test. A sample with four tags can hide an error in the index arithmetic. The random comparison against the knapsack loop found no mismatch here, and it is the check that would catch an index error. Keep a brute force for every new transform.
Checklist
- Write the counting question as a polynomial: each choice is a factor, and the exponent is the total.
- Check the modulus. For NTT, use 998244353 and check that the length divides 2^23.
- Size the arrays to the full product length, a + b - 1, rounded up to a power of two.
- Divide by the length after the inverse transform.
- Reduce modulo p after every multiplication.
- Compare the NTT product with the quadratic product on random inputs of several lengths.
- Compare the counts with a knapsack loop on random small inputs.
- Print the total count of subsets as a check on the coefficient sum, which must be 2 to the power n when there are no modulus effects.
Exercise
Implement the NTT yourself, without copying the code above. Use it to count subset sums by value for the problem in the worked section. Then test it against the quadratic knapsack loop on random inputs with n from 1 to 30 and prices from 1 to 20, at least 1000 cases. Also test the NTT product against the quadratic product on random polynomials. Print the number of cases and the number of mismatches. Use the brute force as a separate function, so the test does not compare the code with itself.
Prompt to use with an AI tool, if you want help:
Write C++17 code that implements an iterative number theoretic transform modulo
998244353 with primitive root 3, and multiplies two polynomials with it. Then
write a function that counts subsets of n prices by total, where each price is
between 1 and 100, by grouping equal prices and multiplying the factors
sum over j of C(m, j) x^(v j) with divide and conquer. Also write a separate
knapsack loop that does the same count one item at a time. Compare the two on
1000 random inputs with n up to 30 and prices up to 20, and print the number of
mismatches. Use only the standard library. Show the product of 1 + x and 1 + x
and check it is 1 + 2x + x^2.A good answer uses 998244353 and not 10^9 + 7, divides by n after the inverse transform, sizes the array to a + b - 1 rounded up to a power of two, and keeps the brute force independent of the NTT. Check its output on the small product: (1 + x)(1 + x) must print 1 2 1. If it prints 1 2 1 but the subset counts fail on random inputs, look at the padding, and if the NTT check fails only at some lengths, look at the bit reversal.
Your own counts are the answer key. Use the sample in this lesson, 1 2 2 3 3 2 2 1, as a first check.
Practice problems
Problem 1. Expand (1 + x + x^2)^2. What is the coefficient of x^2?
Hint: Multiply the polynomial by itself, one term at a time, and collect the x^2 terms.
Show the answer
- The full expansion is 1 + 2x + 3x^2 + 2x^3 + x^4. The x^2 terms come from 1 times x^2, x times x, and x^2 times 1.
Problem 2. How many ways are there to roll a total of 10 with three fair dice, faces 1 to 6?
Hint: This is the coefficient of x^10 in (x + x^2 + ... + x^6)^3. You can also count ordered triples with sum 10.
Show the answer
- We checked this by counting all 216 ordered triples. Of these, 27 have sum 10.
Problem 3. Why does the NTT modulo 998244353 work for lengths up to 2^23, while modulo 1000000007 it works only for length 2?
Hint: Factor p - 1 for each prime. The power of two in p - 1 sets the longest length.
Show the answer
998244352 = 2^23 times 7 times 17, so lengths up to 2^23 have roots. 1000000006 = 2 times 500000003, and 500000003 is prime, so only the length 2 has a root. Lengths 4 and up do not.
Problem 4. The set has four elements. How many of its subsets have total 5?
Hint: Write the generating function as (1 + x)(1 + x^2)(1 + x^4)(1 + x^8). Think about which subsets of powers of two sum to 5, and whether any other subset can do the same.
Show the answer
- Each number from 0 to 15 is a sum of distinct powers of two in exactly one way, so the coefficient of x^5 is 1. The subset is .
Summary
A counting question becomes a polynomial. Each choice is a factor, the exponent records the total, and the coefficient records the count. Multiplying the factors is a convolution, and the NTT does it in about n log n time when the modulus is 998244353 and the length is a power of two up to 2^23. The two checks that catch most mistakes are a random comparison of the NTT with the quadratic product, and a random comparison of the counts with a knapsack loop.