Skip to content
Embercode
✕
CombinatoricsLesson 11 of 11

50 min read

Combinatorics problem set

Why this matters when you build with AI

An AI assistant can write a counting program in seconds, and the program often prints the right answer on the sample. The failures show up later: a factorial that is zero modulo the prime, a missing minus sign in an inclusion-exclusion sum, an answer that is negative, or an index that goes below zero when k is larger than n. The sample is one input. It does not test the cases that break the code.

A brute force is a slow program that is easy to trust, because it does the counting the direct way on small inputs. A stress test runs the fast program and the brute force on many random small inputs and stops at the first difference. Together they catch most of the failures above in a few minutes. Use them on every program an AI writes for you, and on every program you write yourself.

Here is a prompt you can copy. It asks for the brute force and the test loop, not only the solution.

Write a C++17 solution for the problem below. Use the modulus 1000000007. Then write a Python brute force that counts the answer directly for n up to 8, and a generator of random small inputs. Do not say the solution is correct until I have run both programs on at least 200 random inputs. List every input where the answer is 0 and explain why it is 0.

A good answer has three files, a solution that compiles with g++ -O2 -std=c++17, a brute force that does not reuse the formula from the solution, and a note on the cases where k is larger than n or the answer is 0. Check it by running the loop in the next section yourself. A pass on the sample is not a test.

How to use this set

The set has fourteen problems in rising order of difficulty. Problem 1 is solved in full below so you can see the format of a solution and of its check. For problems 2 to 14, try the statement first. Give yourself a time box, for example 15 minutes for the early problems and 45 minutes for the last ones, then open the answer.

For each problem, do three things in order. First write the brute force. Second write the fast solution. Third run the stress loop with a generator you wrote. Do not copy a generator or test list from a site. The statements here are original, and the random tests used to check the answers in this course were generated by scripts, not taken from any judge.

The answers below were checked on one machine (Ubuntu on WSL2, g++ 15.2.0, -O2 -std=c++17). Each sample output shown is the real output of the program. A brute force only covers small inputs, so limits such as 10^18 are checked by reasoning and by the fast solution's own tests, not by brute force. Each answer also says how many random inputs it was checked on, and the result. Timings are for the largest inputs I generated on that machine. Your times will differ.

The stress loop

This loop compares a solution with a brute force on seeds 1 to 200. The solution is sol, the generator is gen.py and the brute force is brute.py. Write these three files yourself. The loop stops at the first seed where the two outputs differ and prints nothing else when all seeds agree.

for i in $(seq 1 200); do
  python3 gen.py $i > in.txt
  ./sol < in.txt > out_sol.txt
  python3 brute.py < in.txt > out_brute.txt
  cmp -s out_sol.txt out_brute.txt || { echo "differ on seed $i"; break; }
done
echo "loop finished"

I ran this loop on problem 1 with a generator and a brute force written for it, over 50 seeds. It printed only loop finished. Seeds are numbers, so a failure can be replayed by running the generator with the same seed.

Problem 1: Pascal table with a composite modulus (worked)

Problem 1 asks for binomial coefficients modulo 1000000000, which is not a prime. Statement: answer q queries. Each query gives n and k. Print C(n,k) modulo 1000000000, where C(n,k) is the number of ways to choose k items from n. If k is less than 0 or greater than n, print 0.

Input: the first line holds q. Each of the next q lines holds n and k.

Limits: 1 <= q <= 1000, 0 <= n <= 1000, -1 <= k <= 1001.

Sample input:

4
5 2
10 5
4 7
1000 500

Sample output:

10
252
0
821216320

The idea is that the factorial method from the first lesson does not work here. It divides by factorials, and dividing modulo a number needs a modular inverse, which exists for a prime modulus. The number 1000000000 is 2 to the power 9 times 5 to the power 9, so it is not prime. The safe method is Pascal's rule, which only adds numbers: C(n,k) equals C(n-1,k-1) plus C(n-1,k). Because the rule only adds, it works for any modulus.

With at most 1000 values of n, the table of all C(n,k) for n up to 1000 fits in memory: about 1001 by 1001 cells of 8 bytes each, which is about 8 MB. Each cell is filled with one addition and one reduction.

The solution below is the program I compiled and ran. Each file starts with #include <bits/stdc++.h> and using namespace std;, which the code below leaves out.

const long long MOD = 1000000000LL;
const int N = 1000;
long long C[N + 1][N + 1];
 
int main() {
    for (int n = 0; n <= N; n++) {
        C[n][0] = 1;
        for (int k = 1; k <= n; k++) C[n][k] = (C[n - 1][k - 1] + C[n - 1][k]) % MOD;
    }
    int q;
    scanf("%d", &q);
    while (q--) {
        long long n, k;
        scanf("%lld %lld", &n, &k);
        long long ans = (k < 0 || k > n) ? 0 : C[n][k];
        printf("%lld\n", ans);
    }
}

Its output on the sample is the four lines shown above: 10, 252, 0, 821216320. Each query is answered with a table read, so the time is the table build, about 10^6 additions, plus q reads.

The brute force for this problem uses Python's exact math.comb, which gives the true value, then takes the remainder. I checked the solution on 80 random inputs with n up to 1000 and k from -1 to 1001, and all 80 matched.

Pitfalls for this problem: C[n-1][k] is read when k equals n. That cell must be zero, which is why the table is a global array, zero at the start. Without the check on k < 0 and k > n, the read happens outside the row, which is undefined behaviour in C++. The output may be a wrong number or a crash.

Problems 2 to 14

Problem 2: Binomials modulo a prime with factorials

Statement: answer q queries. Each query gives n and k. Print C(n,k) modulo 998244353. If k is less than 0 or greater than n, print 0.

Input: the first line holds q. Each of the next q lines holds n and k.

Limits: 1 <= q <= 100000, 0 <= n <= 2000000, -1 <= k <= 2000001.

Sample input:

3
5 2
0 0
3 -1

Sample output:

10
1
0

Show the answer

Idea: precompute factorials and inverse factorials modulo the prime, then C(n,k) is f[n] * inv[k] * inv[n-k]. The modulus is larger than n, so no factorial in the table is zero.

const long long MOD = 998244353LL;
const int MX = 2000000;
 
long long pw(long long b, long long e) {
    long long r = 1;
    b %= MOD;
    while (e) {
        if (e & 1) r = r * b % MOD;
        b = b * b % MOD;
        e >>= 1;
    }
    return r;
}
 
int main() {
    vector<long long> f(MX + 1), inv(MX + 1);
    f[0] = 1;
    for (int i = 1; i <= MX; i++) f[i] = f[i - 1] * i % MOD;
    inv[MX] = pw(f[MX], MOD - 2);
    for (int i = MX; i >= 1; i--) inv[i - 1] = inv[i] * i % MOD;
    int q;
    scanf("%d", &q);
    while (q--) {
        long long n, k;
        scanf("%lld %lld", &n, &k);
        long long ans = (k < 0 || k > n) ? 0 : f[n] * inv[k] % MOD * inv[n - k] % MOD;
        printf("%lld\n", ans);
    }
}

Brute force: math.comb(n, k) % 998244353 for n up to 300. Checked: sample matches; 80 random inputs, 0 mismatches. Complexity: O(2000000 + q) time, about 32 MB for the two tables.

Problem 3: Lucas' theorem for huge n

Statement: answer q queries. Each query gives n and k, with n up to 10^18. Print C(n,k) modulo 7. If k is greater than n, print 0.

Input: the first line holds q. Each of the next q lines holds n and k.

Limits: 1 <= q <= 100000, 0 <= n <= 1000000000000000000, 0 <= k <= 1000000000000000000.

Sample input:

3
10 3
100 50
1000000000000000000 1

Sample output:

1
4
1

Show the answer

Idea: write n and k in base 7. Lucas' theorem says C(n,k) modulo 7 is the product of C(n_i, k_i) modulo 7 over the base 7 digits. Each digit pair is at most 6 by 6, so a 7 by 7 table covers all cases. If some digit of k is larger than the digit of n at the same place, the answer is 0.

int T[7][7];
 
int main() {
    for (int n = 0; n < 7; n++) {
        T[n][0] = 1;
        for (int k = 1; k <= n; k++) T[n][k] = (T[n - 1][k - 1] + T[n - 1][k]) % 7;
    }
    int q;
    scanf("%d", &q);
    while (q--) {
        long long n, k;
        scanf("%lld %lld", &n, &k);
        if (k > n) {
            puts("0");
            continue;
        }
        long long res = 1;
        while (n > 0 || k > 0) {
            int a = n % 7, b = k % 7;
            if (b > a) {
                res = 0;
                break;
            }
            res = res * T[a][b] % 7;
            n /= 7;
            k /= 7;
        }
        printf("%lld\n", res);
    }
}

Brute force: math.comb(n, k) % 7 for n up to 300, where the exact value is known. Checked: sample matches; 80 random inputs, 0 mismatches. Complexity: O(log_7 n) per query, about 22 digit steps for 10^18.

Problem 4: Paths that stay below the diagonal

Statement: a path starts at (0,0). Each step moves one unit right (x grows by 1) or one unit up (y grows by 1). Given a and b with b <= a, count the paths from (0,0) to (a,b) where every point (x,y) on the path has y <= x. Print the count modulo 1000000007.

Input: the first line holds q. Each of the next q lines holds a and b.

Limits: 1 <= q <= 100000, 0 <= b <= a <= 1000000.

Sample input:

3
1 1
2 1
3 3

Sample output:

1
2
5

Check the sample by hand for (2,1): the three paths are RRU, RUR and URR. URR goes through (0,1), where y is larger than x, so it is not allowed. Two paths remain.

Show the answer

Idea: count all paths, then subtract the bad ones. A bad path touches the line y = x + 1. Reflect the part of the path before its first touch across that line. The reflected paths end at (b-1, a+1), so there are C(a+b, b-1) bad paths. The answer is C(a+b, b) minus C(a+b, b-1).

const long long MOD = 1000000007LL;
const int MX = 2000005;
 
long long pw(long long b, long long e) {
    long long r = 1;
    b %= MOD;
    while (e) {
        if (e & 1) r = r * b % MOD;
        b = b * b % MOD;
        e >>= 1;
    }
    return r;
}
 
int main() {
    vector<long long> f(MX + 1), inv(MX + 1);
    f[0] = 1;
    for (int i = 1; i <= MX; i++) f[i] = f[i - 1] * i % MOD;
    inv[MX] = pw(f[MX], MOD - 2);
    for (int i = MX; i >= 1; i--) inv[i - 1] = inv[i] * i % MOD;
    auto C = [&](long long n, long long k) -> long long {
        if (k < 0 || k > n || n < 0) return 0;
        return f[n] * inv[k] % MOD * inv[n - k] % MOD;
    };
    int q;
    scanf("%d", &q);
    while (q--) {
        long long a, b;
        scanf("%lld %lld", &a, &b);
        long long ans = (C(a + b, b) - C(a + b, b - 1) + MOD) % MOD;
        printf("%lld\n", ans);
    }
}

Brute force: a grid dynamic programme that sets dp to 0 for every cell with y > x. Checked: sample matches; 80 random inputs with a up to 10, 0 mismatches.

Problem 5: Permutations with forbidden pairs

Statement: a permutation p of the numbers 1 to n puts each value exactly once. You are given m forbidden pairs (i, v). The permutation must not have p_i equal to v. The positions i are distinct. Values v may repeat. Count the permutations that satisfy every pair, modulo 1000000007.

Input: the first line holds n and m. Each of the next m lines holds i and v.

Limits: 1 <= n <= 200000, 0 <= m <= min(n, 20), 1 <= i <= n, 1 <= v <= n.

Sample input:

3 2
1 1
2 2

Sample output:

3

Check by hand: the six permutations of 1, 2, 3 are listed by their first two entries. 123 and 132 have p1 = 1. 321 has p2 = 2. The other three, 213, 231 and 312, are valid, so the answer is 3.

Show the answer

Idea: inclusion-exclusion over the subsets of the forbidden pairs. For a subset S, count the permutations that break every pair in S, and add or subtract by the size of S. If two pairs in S share a value, no permutation fits, so that subset counts 0. Otherwise the pairs fix |S| positions, leaving (n - |S|)! permutations. Over all 2^m subsets, the sum is the answer.

typedef long long ll;
const ll MOD = 1000000007LL;
 
int main() {
    int n, m;
    scanf("%d %d", &n, &m);
    vector<int> val(m);
    for (int i = 0; i < m; i++) {
        int pos, v;
        scanf("%d %d", &pos, &v);
        val[i] = v;
    }
    vector<unsigned> conf(m, 0);
    for (int i = 0; i < m; i++)
        for (int j = 0; j < m; j++)
            if (i != j && val[i] == val[j]) conf[i] |= 1u << j;
    vector<ll> f(n + 1);
    f[0] = 1;
    for (int i = 1; i <= n; i++) f[i] = f[i - 1] * i % MOD;
    ll ans = 0;
    for (unsigned mask = 0; mask < (1u << m); mask++) {
        bool ok = true;
        int cnt = 0;
        for (int i = 0; i < m && ok; i++) {
            if (mask >> i & 1) {
                cnt++;
                if (conf[i] & mask) ok = false;
            }
        }
        if (!ok) continue;
        ll term = f[n - cnt];
        if (cnt & 1) ans = (ans - term + MOD) % MOD;
        else ans = (ans + term) % MOD;
    }
    printf("%lld\n", ans);
}

Brute force: itertools.permutations over 1 to n for n up to 7, testing each pair. Checked: sample matches; 80 random inputs, 0 mismatches. Complexity: O(2^m * m) time, about 2 times 10^7 steps at m = 20, and O(n) memory.

Problem 6: Bounded compositions

Statement: count the vectors (x_1, ..., x_k) of whole numbers with 0 <= x_i <= u for every i, and with x_1 + ... + x_k = n. Print the count modulo 1000000007.

Input: one line holds k, n and u.

Limits: 1 <= k <= 30, 0 <= n <= 100000, 0 <= u <= 100000.

Sample input:

2 3 2

Sample output:

2

The two vectors are (1,2) and (2,1).

Show the answer

Idea: start from the count of all vectors with sum n, which is C(n + k - 1, k - 1) by stars and bars. Subtract the vectors where some x_i is at least u + 1. Inclusion-exclusion over the set of such variables gives the sum below. Each forced variable uses up u + 1 units of the sum.

typedef long long ll;
const ll MOD = 1000000007LL;
 
ll pw(ll b, ll e) {
    ll r = 1;
    b %= MOD;
    while (e) {
        if (e & 1) r = r * b % MOD;
        b = b * b % MOD;
        e >>= 1;
    }
    return r;
}
 
int main() {
    ll k, n, u;
    scanf("%lld %lld %lld", &k, &n, &u);
    int MX = n + k + 5;
    vector<ll> f(MX + 1), inv(MX + 1);
    f[0] = 1;
    for (int i = 1; i <= MX; i++) f[i] = f[i - 1] * i % MOD;
    inv[MX] = pw(f[MX], MOD - 2);
    for (int i = MX; i >= 1; i--) inv[i - 1] = inv[i] * i % MOD;
    auto C = [&](ll a, ll b) -> ll {
        if (b < 0 || a < 0 || b > a) return 0;
        return f[a] * inv[b] % MOD * inv[a - b] % MOD;
    };
    ll ans = 0;
    for (ll j = 0; j <= k; j++) {
        ll rem = n - j * (u + 1);
        if (rem < 0) break;
        ll term = C(k, j) * C(rem + k - 1, k - 1) % MOD;
        if (j & 1) ans = (ans - term + MOD) % MOD;
        else ans = (ans + term) % MOD;
    }
    printf("%lld\n", ans);
}

Brute force: loop over product(range(u + 1), repeat=k) and count the vectors with sum n, for k up to 4 and u up to 12. Checked: sample matches; 80 random inputs, 0 mismatches. Complexity: O(n + k) for the tables, O(k) for the sum.

Problem 7: Stirling and Bell numbers

Statement: the Stirling number of the second kind S(n,k) is the number of ways to split n labelled items into k non-empty groups, where the groups are not ordered. The Bell number B_n is the sum of S(n,k) over all k, the number of ways to split n items into any number of groups. Each query is one of two types. Type 1 gives n and k and asks for S(n,k), which is 0 when k is greater than n. Type 2 gives n and asks for B_n. Print both modulo 1000000007.

Input: the first line holds q. Each query line is 1 n k for type 1, or 2 n for type 2.

Limits: 1 <= q <= 100000, 0 <= n <= 2000, 0 <= k <= 2000.

Sample input:

4
1 4 2
1 5 3
2 4
1 0 0

Sample output:

7
25
15
1

Show the answer

Idea: S(n,k) = k * S(n-1,k) + S(n-1,k-1). Item n either joins one of the k groups of a split of the first n-1 items, which gives k choices, or it starts a new group alone, which gives S(n-1,k-1). Fill the whole table once, then B_n is a row sum.

typedef long long ll;
const ll MOD = 1000000007LL;
const int N = 2000;
static int S[N + 1][N + 1];
 
int main() {
    S[0][0] = 1;
    for (int n = 1; n <= N; n++)
        for (int k = 1; k <= n; k++)
            S[n][k] = (int)(((ll)k * S[n - 1][k] + S[n - 1][k - 1]) % MOD);
    int q;
    scanf("%d", &q);
    while (q--) {
        int t;
        scanf("%d", &t);
        if (t == 1) {
            int n, k;
            scanf("%d %d", &n, &k);
            printf("%d\n", S[n][k]);
        } else {
            int n;
            scanf("%d", &n);
            ll b = 0;
            for (int k = 0; k <= n; k++) b = (b + S[n][k]) % MOD;
            printf("%lld\n", b);
        }
    }
}

Brute force: enumerate restricted growth strings, which are the labels of the groups in order of first use, and count them by the number of groups. Checked: sample matches; 80 random inputs with n up to 7, 0 mismatches. Complexity: O(N^2) time, about 4 million steps, and about 16 MB for the table. The worst case is 100000 type 2 queries at O(n) each, which ran in 0.19 s on my machine.

Problem 8: Bracket sequences with a height cap

Statement: a bracket sequence of length 2n uses ( and ). It is balanced when the running balance, which is the number of ( so far minus the number of ) so far, never goes below 0 and ends at 0. Count the balanced sequences of length 2n whose running balance never goes above h. Print the count modulo 1000000007.

Input: the first line holds t. Each of the next t lines holds n and h.

Limits: 1 <= t <= 5, 1 <= n <= 2000, 0 <= h <= n.

Sample input:

2
3 1
3 3

Sample output:

1
5

Check by hand: with n = 3 and h = 1, the only balance path that stays in 0 to 1 is ()()(). With h = 3 the count is 5, the Catalan number C_3.

Show the answer

Idea: process the sequence one bracket at a time. Keep an array dp[b] that counts the prefixes whose balance is b. Each step moves every balance up or down by one, dropping any move that leaves the range 0 to h. After 2n steps, dp[0] is the answer.

typedef long long ll;
const ll MOD = 1000000007LL;
 
int main() {
    int t;
    scanf("%d", &t);
    while (t--) {
        int n, h;
        scanf("%d %d", &n, &h);
        vector<ll> dp(h + 2, 0), nd(h + 2, 0);
        dp[0] = 1;
        for (int step = 0; step < 2 * n; step++) {
            fill(nd.begin(), nd.end(), 0);
            for (int b = 0; b <= h; b++) {
                if (!dp[b]) continue;
                if (b + 1 <= h) nd[b + 1] = (nd[b + 1] + dp[b]) % MOD;
                if (b >= 1) nd[b - 1] = (nd[b - 1] + dp[b]) % MOD;
            }
            swap(dp, nd);
        }
        printf("%lld\n", dp[0]);
    }
}

Brute force: itertools.product over the strings of length 2n made of ( and ), checking the balance rule. Checked: sample matches; 40 random inputs with n up to 6, 0 mismatches. Complexity: O(n * h) per test, about 8 million steps at n = h = 2000.

Problem 9: Surjections

Statement: count the functions f from the set {1, ..., n} to the set {1, ..., k} that hit every value, meaning every value in {1, ..., k} is f(x) for at least one x. Print each count modulo 1000000007.

Input: the first line holds q. Each of the next q lines holds n and k.

Limits: 1 <= q <= 10, 1 <= n <= 200000, 1 <= k <= 200000.

Sample input:

3
3 2
4 4
2 3

Sample output:

6
24
0

Check: for n = 3 and k = 2, the 2^3 = 8 functions include 2 that map every input to one value, so 8 - 2 = 6. For n = 2 and k = 3 there are not enough inputs to hit three values, so the answer is 0.

Show the answer

Idea: inclusion-exclusion over the values that are missed. Choosing j values to miss leaves k - j values, and each input has k - j choices. The sum is sum over j of (-1)^j * C(k,j) * (k-j)^n.

typedef long long ll;
const ll MOD = 1000000007LL;
const int MX = 200005;
 
ll pw(ll b, ll e) {
    ll r = 1;
    b %= MOD;
    while (e) {
        if (e & 1) r = r * b % MOD;
        b = b * b % MOD;
        e >>= 1;
    }
    return r;
}
 
int main() {
    vector<ll> f(MX + 1), inv(MX + 1);
    f[0] = 1;
    for (int i = 1; i <= MX; i++) f[i] = f[i - 1] * i % MOD;
    inv[MX] = pw(f[MX], MOD - 2);
    for (int i = MX; i >= 1; i--) inv[i - 1] = inv[i] * i % MOD;
    auto C = [&](ll a, ll b) -> ll {
        if (b < 0 || b > a) return 0;
        return f[a] * inv[b] % MOD * inv[a - b] % MOD;
    };
    int q;
    scanf("%d", &q);
    while (q--) {
        ll n, k;
        scanf("%lld %lld", &n, &k);
        ll ans = 0;
        for (ll j = 0; j <= k; j++) {
            ll term = C(k, j) * pw(k - j, n) % MOD;
            if (j & 1) ans = (ans - term + MOD) % MOD;
            else ans = (ans + term) % MOD;
        }
        printf("%lld\n", ans);
    }
}

Brute force: itertools.product(range(k), repeat=n), keeping the tuples with len(set(f)) == k, for n up to 6 and k up to 7. Checked: sample matches; 80 random inputs, 0 mismatches. Complexity: O(k log n) per query.

Problem 10: Bracelets

Statement: a bracelet is a ring of n beads, and each bead gets one of k colours. Two colourings are the same bracelet when one can be turned into the other by rotating the ring or by flipping it over. Count the bracelets modulo 1000000007.

Input: the first line holds q. Each of the next q lines holds n and k.

Limits: 1 <= q <= 1000, 1 <= n <= 1000000, 1 <= k <= 1000000000.

Sample input:

4
1 1
2 2
3 2
4 3

Sample output:

1
3
4
21

Warning: counting only rotations (necklaces) is a different answer. For n = 6 and k = 2 there are 14 necklaces and 13 bracelets. A program that skips the flips prints 14 where 13 is correct.

Show the answer

Idea: Burnside's lemma. The number of distinct objects equals the average number of colourings fixed by each symmetry. Here the symmetries are the 2n elements of the dihedral group, the rotations and the flips. A rotation whose cycles have length d, for each divisor d of n, leaves n/d cycles of beads, so it fixes k^(n/d) colourings. There are phi(d) such rotations. A flip fixes colourings that are constant on each cycle of the flip. For odd n every flip has one fixed bead and (n-1)/2 pairs, so it has (n+1)/2 cycles. For even n there are n/2 flips through two beads, with n/2 + 1 cycles, and n/2 flips through edges, with n/2 cycles. The total is divided by 2n.

typedef long long ll;
const ll MOD = 1000000007LL;
 
ll pw(ll b, ll e) {
    ll r = 1;
    b %= MOD;
    while (e) {
        if (e & 1) r = r * b % MOD;
        b = b * b % MOD;
        e >>= 1;
    }
    return r;
}
 
int main() {
    int q;
    scanf("%d", &q);
    while (q--) {
        ll n, k;
        scanf("%lld %lld", &n, &k);
        vector<pair<ll, int>> fac;
        ll m = n;
        for (ll p = 2; p * p <= m; p++) {
            if (m % p == 0) {
                int e = 0;
                while (m % p == 0) {
                    m /= p;
                    e++;
                }
                fac.push_back({p, e});
            }
        }
        if (m > 1) fac.push_back({m, 1});
        vector<pair<ll, ll>> cur = {{1, 1}};
        for (auto [p, e] : fac) {
            vector<pair<ll, ll>> nxt;
            for (auto [d, ph] : cur) {
                nxt.push_back({d, ph});
                ll pk = 1;
                for (int a = 1; a <= e; a++) {
                    pk *= p;
                    nxt.push_back({d * pk, ph * (pk / p) * (p - 1)});
                }
            }
            cur = nxt;
        }
        ll rot = 0;
        for (auto [d, ph] : cur) rot = (rot + ph % MOD * pw(k, n / d)) % MOD;
        ll ref;
        if (n % 2 == 1) {
            ref = (n % MOD) * pw(k, (n + 1) / 2) % MOD;
        } else {
            ll h = n / 2;
            ref = (h % MOD) * ((pw(k, h + 1) + pw(k, h)) % MOD) % MOD;
        }
        ll total = (rot + ref) % MOD * pw(2 * n % MOD, MOD - 2) % MOD;
        printf("%lld\n", total);
    }
}

The code loops over the divisors d of n through the factorisation, which gives both the divisor and phi(d). Brute force: for each colouring, take the smallest string among its 2n rotations and reversed rotations, and count the distinct smallest strings, for n up to 7 and k up to 4. Checked: sample matches; 25 random inputs, 0 mismatches. Complexity: O(sqrt(n) + d(n) log k) per query, where d(n) is the number of divisors.

Problem 11: Grid paths around blocked cells

Statement: a grid has n rows and m columns. Rows are numbered 1 to n from the top and columns 1 to m from the left. A path starts at (1,1) and ends at (n,m). Each step moves one cell right or one cell down. Some k cells are blocked, and a path may not enter a blocked cell. Count the paths modulo 1000000007. The start and the end are not blocked.

Input: the first line holds n, m and k. Each of the next k lines holds a row r and a column c.

Limits: 1 <= n, m <= 1000000, 0 <= k <= 2000. The blocked cells are distinct and are not (1,1) or (n,m).

Sample input:

2 3 1
1 2

Sample output:

1

Check: a 2 by 3 grid has C(3,1) = 3 paths. The blocked cell (1,2) is on two of them, the ones that start with a right step. One path is left.

Show the answer

Idea: sort the blocked cells. Let dp[i] be the number of paths from the start to blocked cell i that do not touch any earlier blocked cell. A path to cell i is counted in the total of all paths, minus the paths that first touch an earlier blocked cell j and then continue to i. The answer is the total to the end, minus paths through each blocked cell. Paths between two cells are binomial coefficients, with a zero count when the target is not down-right of the source.

typedef long long ll;
const ll MOD = 1000000007LL;
 
ll pw(ll b, ll e) {
    ll r = 1;
    b %= MOD;
    while (e) {
        if (e & 1) r = r * b % MOD;
        b = b * b % MOD;
        e >>= 1;
    }
    return r;
}
 
int main() {
    ll n, m;
    int k;
    scanf("%lld %lld %d", &n, &m, &k);
    vector<pair<ll, ll>> b(k);
    for (auto& x : b) scanf("%lld %lld", &x.first, &x.second);
    sort(b.begin(), b.end());
    int MX = n + m + 5;
    vector<ll> f(MX + 1), inv(MX + 1);
    f[0] = 1;
    for (int i = 1; i <= MX; i++) f[i] = f[i - 1] * i % MOD;
    inv[MX] = pw(f[MX], MOD - 2);
    for (int i = MX; i >= 1; i--) inv[i - 1] = inv[i] * i % MOD;
    auto C = [&](ll a, ll c) -> ll {
        if (c < 0 || c > a) return 0;
        return f[a] * inv[c] % MOD * inv[a - c] % MOD;
    };
    auto paths = [&](ll r1, ll c1, ll r2, ll c2) -> ll {
        ll dr = r2 - r1, dc = c2 - c1;
        if (dr < 0 || dc < 0) return 0;
        return C(dr + dc, dr);
    };
    vector<ll> dp(k, 0);
    for (int i = 0; i < k; i++) {
        dp[i] = paths(1, 1, b[i].first, b[i].second);
        for (int j = 0; j < i; j++) {
            ll through = dp[j] * paths(b[j].first, b[j].second, b[i].first, b[i].second) % MOD;
            dp[i] = (dp[i] - through + MOD) % MOD;
        }
    }
    ll ans = paths(1, 1, n, m);
    for (int j = 0; j < k; j++) {
        ll through = dp[j] * paths(b[j].first, b[j].second, n, m) % MOD;
        ans = (ans - through + MOD) % MOD;
    }
    printf("%lld\n", ans);
}

Sorting by (row, column) guarantees that any blocked cell a path can reach before cell i comes earlier in the list. Brute force: a grid dynamic programme where blocked cells keep the value 0, for grids up to 8 by 8 with up to 6 blocked cells. Checked: sample matches; 80 random inputs, 0 mismatches. Complexity: O(k^2 + n + m) time, with the factorial table of size n + m. My largest test (10^6 by 10^6 with 2000 blocked cells) ran in 0.12 s.

Problem 12: Strings with forbidden neighbours by matrix power

Statement: a string of length n uses the symbols 0, 1 and 2. The forbidden neighbour pairs are "01", "10" and "22". Count the strings of length n that contain none of these pairs as two adjacent symbols. Print the count modulo 1000000007.

Input: the first line holds q. Each of the next q lines holds one n.

Limits: 1 <= q <= 100000, 1 <= n <= 1000000000000000000.

Sample input:

3
1
2
3

Sample output:

3
6
12

Check: length 1 has 3 strings. Length 2 has 9 strings, and 3 of them are forbidden, so 6 remain.

Show the answer

Idea: keep one count per last symbol. A matrix M has M[a][b] = 1 when the pair (a, b) is allowed. The number of strings of length n is the sum of all entries of M^(n-1), because the first symbol is free. Fast power of the matrix takes O(log n) multiplications of 3 by 3 matrices.

typedef long long ll;
const ll MOD = 1000000007LL;
typedef array<array<ll, 3>, 3> M3;
 
M3 mul(const M3& A, const M3& B) {
    M3 C{};
    for (int i = 0; i < 3; i++)
        for (int j = 0; j < 3; j++)
            for (int k = 0; k < 3; k++) C[i][j] = (C[i][j] + A[i][k] * B[k][j]) % MOD;
    return C;
}
 
int main() {
    M3 base{}, id{};
    for (int i = 0; i < 3; i++) {
        id[i][i] = 1;
        for (int j = 0; j < 3; j++) {
            bool forbidden = (i == 0 && j == 1) || (i == 1 && j == 0) || (i == 2 && j == 2);
            base[i][j] = forbidden ? 0 : 1;
        }
    }
    int q;
    scanf("%d", &q);
    while (q--) {
        ll n;
        scanf("%lld", &n);
        ll e = n - 1;
        M3 P = id, B = base;
        while (e > 0) {
            if (e & 1) P = mul(P, B);
            B = mul(B, B);
            e >>= 1;
        }
        ll ans = 0;
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++) ans = (ans + P[i][j]) % MOD;
        printf("%lld\n", ans);
    }
}

Brute force: a dynamic programme over the last symbol for n up to 2000, which is O(n) and is the check for the matrix. Checked: sample matches; 80 random inputs with n up to 2000, 0 mismatches. The n up to 10^18 limit cannot be checked by brute force; it is covered by the log-time power. My largest test (100000 queries of random n up to 10^18) ran in 0.32 s.

Problem 13: Subset sums with NTT

Statement: you are given n numbers a_1 to a_n. For each query s, count the subsets of the indices whose values add up to exactly s. The empty subset counts for s = 0. Print each count modulo 998244353.

Input: the first line holds n. The second line holds the n values. The next line holds q. Each of the next q lines holds one s.

Limits: 1 <= n <= 200, 1 <= a_i <= 100, 1 <= q <= 1000, 0 <= s <= 20000.

Sample input:

3
1 2 3
2
3
6

Sample output:

2
1

Check: for s = 3 the subsets are {3} and {1,2}. For s = 6 the only subset is .

Show the answer

Idea: the generating polynomial for item a is 1 + x^a. The product of all these polynomials has, as its coefficient of x^s, the number of subsets with sum s. Multiplying polynomials is a convolution, which the number theoretic transform (NTT) does in O(N log N) for a modulus with a suitable root of unity. The modulus 998244353 is 119 * 2^23 + 1 with primitive root 3. The program multiplies the factors in a balanced tree, so the sizes stay close.

typedef long long ll;
const ll MOD = 998244353LL, G = 3;
 
ll pw(ll b, ll e) {
    ll r = 1;
    b %= MOD;
    if (b < 0) b += MOD;
    while (e) {
        if (e & 1) r = r * b % MOD;
        b = b * b % MOD;
        e >>= 1;
    }
    return r;
}
 
void ntt(vector<ll>& a, bool invert) {
    int n = a.size();
    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) {
        ll w = pw(G, (MOD - 1) / len);
        if (invert) w = pw(w, MOD - 2);
        for (int i = 0; i < n; i += len) {
            ll wn = 1;
            for (int j = 0; j < len / 2; j++) {
                ll u = a[i + j], v = a[i + j + len / 2] * wn % MOD;
                a[i + j] = (u + v) % MOD;
                a[i + j + len / 2] = (u - v + MOD) % MOD;
                wn = wn * w % MOD;
            }
        }
    }
    if (invert) {
        ll ni = pw(n, MOD - 2);
        for (auto& x : a) x = x * ni % MOD;
    }
}
 
vector<ll> mul(const vector<ll>& a, const vector<ll>& b) {
    int need = a.size() + b.size() - 1;
    vector<ll> res(need, 0);
    if (min(a.size(), b.size()) <= 32) {
        for (size_t i = 0; i < a.size(); i++)
            for (size_t j = 0; j < b.size(); j++) res[i + j] = (res[i + j] + a[i] * b[j]) % MOD;
        return res;
    }
    int sz = 1;
    while (sz < need) sz <<= 1;
    vector<ll> fa(a), fb(b);
    fa.resize(sz);
    fb.resize(sz);
    ntt(fa, false);
    ntt(fb, false);
    for (int i = 0; i < sz; i++) fa[i] = fa[i] * fb[i] % MOD;
    ntt(fa, true);
    fa.resize(need);
    return fa;
}
 
vector<ll> prodRange(const vector<int>& a, int l, int r) {
    if (l == r) {
        vector<ll> p(a[l] + 1, 0);
        p[0] = 1;
        p[a[l]] = (p[a[l]] + 1) % MOD;
        return p;
    }
    int mid = (l + r) / 2;
    return mul(prodRange(a, l, mid), prodRange(a, mid + 1, r));
}
 
int main() {
    int n;
    scanf("%d", &n);
    vector<int> a(n);
    for (auto& x : a) scanf("%d", &x);
    vector<ll> P = prodRange(a, 0, n - 1);
    int q;
    scanf("%d", &q);
    while (q--) {
        long long s;
        scanf("%lld", &s);
        printf("%lld\n", (s >= 0 && s < (long long)P.size()) ? P[s] : 0LL);
    }
}

Brute force: the knapsack loop dp[s] += dp[s - a] over the items, for n and values up to 60, which gives sizes large enough to use the NTT path. Checked: sample matches; 80 random inputs, 0 mismatches. Complexity: O(S log^2 S) for the product, with S = sum of the values, at most 20000 here. My largest test (n = 200, values up to 100, 1000 queries) ran in 0.01 s.

Problem 14: Expected rolls for k distinct faces

Statement: a fair die has n faces numbered 1 to n. Roll it until you have seen k different faces. The expected number of rolls is a fraction P/Q. For each query, print P times the modular inverse of Q, modulo 1000000007. The denominators in this problem are products of numbers up to 10^6, which are not multiples of 1000000007, so this value is well defined.

Input: the first line holds q. Each of the next q lines holds n and k.

Limits: 1 <= q <= 10, 1 <= k <= n <= 1000000.

Sample input:

2
2 2
3 2

Sample output:

3
500000006

Check: for n = 2 and k = 2 the expected number of rolls is 3, which is 2/2 + 2/1. For n = 3 and k = 2 it is 3/3 + 3/2 = 5/2, and 500000006 times 2 is 1000000012, which is 5 more than 1000000007, so 500000006 is 5/2 modulo the prime.

Show the answer

Idea: the expected time to go from i seen faces to i + 1 seen faces is n / (n - i), because each roll is new with probability (n - i) / n. Add these times for i from 0 to k - 1. Each term is n times the inverse of n - i, and the inverses come from a table built once with the linear rule inv[i] = -(p / i) * inv[p % i].

typedef long long ll;
const ll MOD = 1000000007LL;
const int MX = 1000000;
 
int main() {
    vector<ll> inv(MX + 1, 1);
    for (int i = 2; i <= MX; i++) inv[i] = (MOD - (MOD / i) * inv[MOD % i] % MOD) % MOD;
    int q;
    scanf("%d", &q);
    while (q--) {
        ll n, k;
        scanf("%lld %lld", &n, &k);
        ll ans = 0;
        for (ll i = 0; i < k; i++) ans = (ans + n % MOD * inv[n - i]) % MOD;
        printf("%lld\n", ans);
    }
}

Brute force: an exact fraction recurrence in Python, Fraction, which solves the equation for the expected value from each state with i seen faces, then converts the fraction to a residue. Checked: sample matches; 80 random inputs with n up to 30, 0 mismatches. Complexity: O(k) per query after one table build of size 10^6.

Common mistakes

Using factorials with a composite modulus, or with n at least as large as the prime. For problem 3 modulo 7, the factorial route gives 0 for C(10,3), because 10! contains the factor 7. The true value is math.comb(10, 3) % 7, which is 1. Use Lucas' theorem, as in problem 3.

Forgetting the subtracted term in the ballot count. For a = 2 and b = 1, math.comb(3, 1) is 3, but the answer is 2, which is math.comb(3, 1) - math.comb(3, 0).

Counting necklaces when the problem asks for bracelets. For n = 6 and k = 2 the rotation-only count is 14 and the bracelet count is 13.

Reading past the end of an array when k is larger than n in a binomial table. In C++ this is undefined behaviour. The result may look plausible, or the program may crash.

Subtracting modulo a prime without adding the modulus back. (a - b) % MOD can be negative in C++, and a negative answer is printed as it is.

Passing the sample and stopping there. The sample has four to five queries. A program that is wrong for k greater than n, or for n equal to 0, may still pass it.

Overflow in the middle of a product. Two residues below 1000000007 multiply to about 10^18, which fits in a 64-bit integer but not in a 32-bit one. Use long long for every product.

Checklist

  • Write the brute force first, and make it count the statement directly.
  • Run the stress loop with at least 200 seeds before you trust a solution.
  • Check k < 0 and k > n before you read any table.
  • Put the modulus after every subtraction, so the result is never negative.
  • Use a modular inverse only when the modulus is prime, and only for values that are not multiples of it.
  • Test n = 0, k = 0, k = n, and the largest n in the limits for speed.
  • Compare the output format with the statement: one answer per line, in query order.

Summary

The fourteen problems use Pascal's rule, factorial tables, Lucas' theorem, reflection, inclusion-exclusion, stars and bars, Stirling and Bell numbers, Burnside's lemma, grid paths with obstacles, matrix power, NTT, and expected values. Each one has a brute force to compare with, and each answer shows the sample output and the number of random inputs it was checked on. The practical habit is the same for all of them: write the brute force, run the stress loop, and read the output for the zero and the negative cases before you submit.

Practice this topic: coding problems for Combinatorics →
← Probability and expected value in contests