[CodeChef] Starters 71 — Expected Sum: the full distribution first, then Vandermonde collapses it
Published 7 October 2026
Problem: CodeChef — Expected Sum (Starters 71, difficulty 2495).
There are \(N\) chits: \(A\) have \(1\) written on them and \(B\) have \(0\) (\(A + B = N\)). Chef and Chefina take turns, Chef first. On each turn the player picks a chit uniformly at random from the ones left, adds its number to their own score, and throws it away. Find Chef’s expected score once every chit is gone, as \(P \cdot Q^{-1} \bmod 998244353\).
Constraints: \(T \le 1000\), \(0 \le A, B \le 10^9\), \(1 \le A + B \le 10^9\).
The one-line answer is \(\dfrac{A \cdot \lceil N/2 \rceil}{N}\), and linearity of expectation gets there almost immediately. I wanted to get there without linearity first: write down the full probability distribution of Chef’s score, take its mean, and see the sum collapse. That route is longer, but it shows why the answer is so simple, and it’s where I made my mistakes.
Setting up¶
Let \(N = A + B\) and let
be the number of turns Chef gets. He moves first, so on an odd \(N\) he gets the extra turn. Chefina gets the other \(N - K\) turns.
I’ll use \(K\) for Chef’s turns from now on. My first notes reused \(a\) and \(b\) for “Chef’s turns” and “Chefina’s turns”, which clashes with \(A\) and \(B\), the chit counts, and made the counting below harder to follow.
Chef’s final score is just the number of \(1\)-chits he picks up. Call it \(S\). Then
(That’s \(s\), not \(s^2\). Weighting by \(s^2\) would give \(\mathbb{E}[S^2]\), the second moment.)
The sample space¶
The whole game is decided by the order in which the chits come out. Write that order as a string of length \(N\) with \(A\) ones and \(B\) zeros; positions \(1, 3, 5, \dots\) go to Chef and positions \(2, 4, 6, \dots\) go to Chefina.
Since every pick is uniform over the remaining chits, every one of these 0/1 strings is equally likely. There are
of them: choose which \(A\) of the \(N\) positions hold a \(1\). The chits with the same number are indistinguishable, so we count strings, not permutations of labelled chits. (Counting labelled permutations works too, but every string just gets multiplied by the same \(A!\,B!\), which cancels.)
So
Counting the favourable strings — where I went wrong¶
My first count was \(\binom{K}{s}\): out of Chef’s \(K\) turns, choose which \(s\) of them get a \(1\), and the rest of his turns get a \(0\). Then I said “the remaining \(A - s\) ones go to Chefina, and we don’t care where they land in her turns.”
That’s the mistake. We don’t care for Chef’s score where they land, but the denominator \(\binom{N}{A}\) counts complete strings, including how Chefina’s turns are filled. The numerator has to count complete strings too, or the ratio isn’t a probability. Each pattern on Chef’s turns extends to several different complete strings, and every one of them is a separate, equally likely outcome.
Chefina’s \(N - K\) turns must hold the remaining \(A - s\) ones, and there are \(\binom{N-K}{A-s}\) ways to place them. So the correct count is
A small example¶
Take \(N = 4\), \(A = 2\), so \(K = 2\) and Chefina also has \(2\) turns. All \(\binom{4}{2} = 6\) strings, with Chef’s positions (1st and 3rd) in bold:
| String | Chef gets | Chefina gets | \(S\) |
|---|---|---|---|
| 1 1 0 0 | 1, 0 | 1, 0 | 1 |
| 1 0 1 0 | 1, 1 | 0, 0 | 2 |
| 1 0 0 1 | 1, 0 | 0, 1 | 1 |
| 0 1 1 0 | 0, 1 | 1, 0 | 1 |
| 0 1 0 1 | 0, 0 | 1, 1 | 0 |
| 0 0 1 1 | 0, 1 | 0, 1 | 1 |
\(S = 1\) shows up \(4\) times, not \(2\). Chef’s turns can read 10 or 01 — that’s
\(\binom{2}{1} = 2\) — and for each of them Chefina’s turns can read 10 or 01 — another
\(\binom{2}{1} = 2\). So \(\Pr[S = 1] = 4/6\), and my \(\binom{K}{s}\)-only count would have said \(2/6\).
Check: \(\Pr[S=0] = 1/6\), \(\Pr[S=1] = 4/6\), \(\Pr[S=2] = 1/6\), which sums to \(1\). The mean is \((0 + 4 + 2)/6 = 1\), and \(A K / N = 2 \cdot 2 / 4 = 1\). ✓
This is the hypergeometric distribution: draw \(K\) items without replacement from a pool of \(N\), of which \(A\) are “good”, and count the good ones you drew.
The range of \(s\)¶
Not every \(s\) is possible:
- Chef can’t get more ones than he has turns or than exist: \(s \le \min(K, A)\).
- Chefina can hold at most \(N - K\) ones, so Chef must take at least the overflow: \(s \ge \max(0,\ A - (N - K))\).
Outside this range one of the binomials is \(0\) anyway (with the convention \(\binom{n}{k} = 0\) for \(k < 0\) or \(k > n\)), so the sum can be written over all \(s\) and nothing changes.
The direct formula, and why it’s too slow¶
Putting it together:
This is correct, and with precomputed factorials it’s one pass over \(s\). But \(A\) can be \(10^9\), so the range of \(s\) can have about \(10^9\) values per test, with \(T = 1000\) tests. And factorials up to \(10^9\) don’t fit in memory. We need the sum in closed form.
Collapsing the sum¶
Two identities do it.
Absorption. Pull the factor \(s\) into the binomial:
Combinatorially: picking a team of \(s\) from \(K\) people and then a captain from the team is the same as picking the captain from all \(K\) first and then the other \(s - 1\) members from the remaining \(K - 1\).
So the numerator becomes
Vandermonde. Look at what’s left. The top numbers add up to a constant, \((K-1) + (N-K) = N - 1\), and so do the bottom numbers, \((s-1) + (A-s) = A - 1\). That’s exactly the shape of Vandermonde’s identity (next section), which says the sum is
Finish. Using \(\binom{N}{A} = \frac{N}{A}\binom{N-1}{A-1}\) (absorption again):
That’s \(O(1)\) arithmetic plus one modular inverse.
Vandermonde’s identity¶
The idea. You have two groups, one with \(m\) items and one with \(n\), and you want to pick \(r\) items in total.
- Right side: pour both groups into one pool of \(m + n\) and choose \(r\). That’s \(\binom{m+n}{r}\) ways.
- Left side: split the same choices by how many items, \(k\), come from the first group. The other \(r - k\) must come from the second. That’s \(\binom{m}{k}\binom{n}{r-k}\) for each \(k\); add over all \(k\).
Both sides count the same selections, so they’re equal. (This is a double-counting proof.)
In our problem: the first group is Chef’s turns other than the captain one (\(K - 1\) of them), the second is Chefina’s turns (\(N - K\)), and we’re placing the \(A - 1\) ones that aren’t the captain. Sum over how many land on Chef’s side, and you just get “place \(A - 1\) ones in \(N - 1\) slots”.
Example 1: a committee. A club has \(5\) men and \(4\) women. How many \(3\)-person committees are there? Directly, \(\binom{9}{3} = 84\). Split by the number of men \(k\):
| Men \(k\) | Women \(3-k\) | Ways |
|---|---|---|
| 0 | 3 | \(\binom{5}{0}\binom{4}{3} = 1 \cdot 4 = 4\) |
| 1 | 2 | \(\binom{5}{1}\binom{4}{2} = 5 \cdot 6 = 30\) |
| 2 | 1 | \(\binom{5}{2}\binom{4}{1} = 10 \cdot 4 = 40\) |
| 3 | 0 | \(\binom{5}{3}\binom{4}{0} = 10 \cdot 1 = 10\) |
\(4 + 30 + 40 + 10 = 84\). ✓
Example 2: sum of squares of a Pascal row. Set \(m = n = r\). Since \(\binom{n}{n-k} = \binom{n}{k}\), the identity becomes
For \(n = 3\), row 3 is \(1, 3, 3, 1\), and \(1 + 9 + 9 + 1 = 20 = \binom{6}{3}\). ✓
Example 3: the polynomial proof. \((1+x)^m (1+x)^n = (1+x)^{m+n}\). The coefficient of \(x^r\) on the right is \(\binom{m+n}{r}\). On the left you get \(x^r\) by taking \(x^k\) from the first factor and \(x^{r-k}\) from the second, for every \(k\), which is \(\sum_k \binom{m}{k}\binom{n}{r-k}\). This is why it’s also called Vandermonde’s convolution: multiplying polynomials convolves their coefficients.
Why it’s worth remembering.
- When you see a sum of products of binomials whose bottom indices add up to a constant, it’s often Vandermonde in disguise. That’s how it showed up here.
- Divide both sides by \(\binom{m+n}{r}\) and each term becomes a hypergeometric probability. So Vandermonde is the statement that the hypergeometric distribution sums to \(1\).
- It generalises to more groups: \(\sum \binom{n_1}{k_1}\cdots\binom{n_p}{k_p} = \binom{n_1 + \dots + n_p}{r}\), summed over \(k_1 + \dots + k_p = r\).
- A grid-path version: a path from \((0,0)\) with \(m + n\) right/up steps crosses the anti-diagonal \(x + y = m\) at exactly one point. Counting paths by where they cross gives Vandermonde’s sum. “Split at a checkpoint” shows up a lot in grid-path counting problems.
The short way: linearity of expectation¶
Now the route I deliberately skipped. Write Chef’s score as a sum over his turns:
The order of the chits is uniformly random, so by symmetry any single position is equally likely to hold any of the \(N\) chits. So each position, looked at on its own, is a \(1\) with probability \(A / N\). That’s the unconditional probability: it doesn’t depend on which position it is. So
The \(X_i\) are not independent (if Chef’s first pick is a \(1\), his second is less likely to be). That doesn’t matter: linearity of expectation holds for dependent variables too. Independence only matters for things like the variance.
The two routes meet: the Vandermonde computation above is exactly linearity of expectation done the long way.
Aside: the variance. For the hypergeometric distribution it’s
The last factor is what dependence costs you compared with \(K\) independent draws.
Modular arithmetic¶
The answer is \(\frac{K \cdot A}{N}\), so print \(K \cdot A \cdot N^{-1} \bmod 998244353\). Since \(998244353\) is prime, \(N^{-1} = N^{p-2} \bmod p\) by Fermat’s little theorem.
The inverse always exists here: \(1 \le N \le 10^9 < 998244353\), so \(N\) is never a multiple of the modulus. No special cases are needed for \(A = 0\) (the product is \(0\)) or \(B = 0\) (the \(A / N\) factor is \(1\), giving \(K\)).
Checking against the samples:
- \(A = 1, B = 1\): \(N = 2, K = 1\), answer \(1/2\), which is \(499122177\). ✓
- \(A = 0, B = 5\): answer \(0\). ✓
- \(A = 3, B = 4\): \(N = 7, K = 4\), answer \(12/7\), which is \(285212674\). ✓
I also checked \(\frac{KA}{N}\) against a brute force over every distinct order of the chits, and against the full-distribution sum, for all \(A, B \le 6\).
Code¶
import sys
MOD = 998244353
def solve():
data = list(map(int, sys.stdin.buffer.read().split()))
t = data[0]
out = []
idx = 1
for _ in range(t):
A = data[idx]
B = data[idx + 1]
idx += 2
N = A + B
K = (N + 1) // 2 # Chef starts, so Chef gets ceil(N / 2) turns
# E[score] = K * A / N; N <= 10^9 < MOD, so N is always invertible
ans = K % MOD * (A % MOD) % MOD
ans = ans * pow(N, MOD - 2, MOD) % MOD
out.append(str(ans))
sys.stdout.write("\n".join(out))
if __name__ == "__main__":
solve()
Complexity: \(O(\log \text{MOD})\) per test case for the modular exponentiation, \(O(T)\) memory for the output.
The full-distribution version, for small inputs¶
This is the formula from before the collapse. It’s useful as a brute-force checker, and it shows the hypergeometric sum is the same thing as \(KA/N\).
from fractions import Fraction
from math import comb
def expected_by_distribution(A, B):
N = A + B
K = (N + 1) // 2
lo = max(0, A - (N - K))
hi = min(K, A)
favourable = sum(s * comb(K, s) * comb(N - K, A - s) for s in range(lo, hi + 1))
return Fraction(favourable, comb(N, A))
assert expected_by_distribution(3, 4) == Fraction(12, 7)
What I took away¶
- The numerator and the denominator must count the same kind of outcome. If the denominator counts complete orders, the numerator has to count complete orders too, even for the part of the order you “don’t care about”. Counting only Chef’s half gave \(\binom{K}{s}\) and dropped the \(\binom{N-K}{A-s}\) factor.
- \(s\binom{K}{s} = K\binom{K-1}{s-1}\), then Vandermonde is the standard way to take the mean of a sum of binomial products. Look for bottom indices that add up to a constant.
- Expected count of “good” things → linearity. Write the count as a sum of indicators, and don’t worry about dependence. Doing it the long way once makes it clear why the shortcut is allowed.