Compositions via dynamic programming

A composition of a nonnegative integer $S$ into $n$ parts is a sequence $(x_1,\ldots,x_n)$ with $x_i \in \mathbb{Z}_{\ge 0}$ and

(If you require every part to be positive, that is a strict composition; the DP below changes slightly. This note uses nonnegative parts, matching stars-and-bars and problems such as E0 in contest 2228.)

We build one core DP in four stages:

1. brute enumeration;
2. straightforward two-table DP;
3. one-row DP with an $O(S)$ transition per variable;
4. binary exponentiation when the same transition is repeated $n$ times.

Throughout, $n \ge 1$ and $S \ge 0$.

The number of compositions is the stars-and-bars count

(for $n \ge 1$). When $S = 0$, the only composition is $(0,\ldots,0)$. Our DP must match this closed form on small tests.

If $n$ and $S$ are tiny, recurse on the first variable: try $x_1 = 0,1,\ldots,S$, recurse on $(x_2,\ldots,x_n)$ with budget $S - x_1$.

count(n, S):
  if n = 0: return 1 if S = 0 else 0
  ans = 0
  for v = 0 .. S:
    ans += count(n - 1, S - v)
  return ans

Time $O(\binom{S+n-1}{n-1})$ in the worst case (proportional to the number of compositions). Fine for checking answers, useless for contest sizes.

See brute.cpp.

Assign variables left to right. Let

Base: $\mathrm{dp}[0][0] = 1$, other $\mathrm{dp}[0][s] = 0$.

Transition (choose $x_i = v$):

Answer: $\mathrm{dp}[n][S]$.

Naively, for each $i$ and $s$ you loop $v$ — $O(n S^2)$ time and $O(nS)$ space.

Rolling the $i$ index gives $O(nS^2)$ time and $O(S)$ space (still the inner $v$ loop).

See dp_two_row.cpp.

The bottleneck is the inner loop over $v$. Fix $i$ and write the update as

So one new variable is a prefix sum of the old row.

Algorithm for one layer:

1. Build prefix sums $\mathrm{pref}[t] = \sum_{s=0}^{t} \mathrm{dp}[s]$.
2. Set $\mathrm{dp}_{\mathrm{new}}[t] = \mathrm{pref}[t]$.

Repeat for $i = 1,\ldots,n$. Time $O(nS)$, space $O(S)$.

dp[0] = 1, others 0
repeat n times:
  pref[t] = sum_{s=0..t} dp[s]
  dp <- pref
answer = dp[S]

This is the same “bulk over $v$” idea used in the E0 / E1 write-ups when a free slot distributes mass among sums $s$.

See dp_prefix_layer.cpp.

When every variable is anonymous (only $n$ and $S$ matter), each layer applies the same map $T$:

We need $T^n$ applied to the initial vector $\mathrm{dp}_0$ with $\mathrm{dp}_0[0]=1$.

Binary exponentiation on the map $T$:

• keep res (composition of maps so far) and cur (current power of $T$);
• if $n$ has a bit set, res ← res ∘ cur;
• cur ← cur ∘ cur;
• answer is res(dp_0)[S].

Each compose costs two $O(S)$ applications of $T$, so $O(S \log n)$ total. Useful when $n$ is large (for example $n \le 3\cdot 10^5$, $S \le 10^6$) and you cannot afford $O(nS)$.

See pow_layers.cpp.

Many problems ask for a sum over compositions, not just a count. Example from E0: for each composition $c$, let $P_i = c_1+\cdots+c_i$ and

You still process variables left to right, but each partial assignment carries extra aggregates ($\sum$ of prefix, $\sum$ of prefix squared, $\sum$ of earlier $P_i^2$, and so on). One free layer stays $O(S)$ with a few prefix arrays; see the E0 editorial for the full transition.

The counting DP in this tutorial is the count-only special case where you only store $\mathrm{dp}[s]$ (our $\mathrm{cnt}[s]$).

• After each stage, compare to $\binom{S+n-1}{n-1}$ on random small $(n,S)$.
• For $S=0$, the answer must be $1$ (all zeros).
• For $n=1$, the answer must be $1$ (only $x_1 = S$).

All programs read $n$, $S$, print the count modulo $998\,244\,353$.

All implementations

• E0. Amanojaku and Sequence (Tutorial) — weighted sum $G(n,m)$ over compositions.
• E1. Amanojaku and Sequence (Easy) — subarray queries; all-free case uses $O(S \log L)$ layering.
• Count arrays with fixed product — stars and bars without DP.