Why this matters when you build with AI
Some contest statements ask for a probability or an expected value, and the answer is a fraction that has to be printed modulo a prime, such as 998244353. An AI tool asked for this often does one of three things. It computes the answer with floating point numbers and prints 0.333333, which is not a fraction and is wrong for the judge. It assumes that the expected value of a product is the product of the expected values, which is false when the parts depend on each other. Or it writes a recurrence over states that loops forever, because it never solves for the state that refers to itself.
These mistakes are easy to catch if you test on small inputs with exact fractions. Python's fractions module keeps numerators and denominators as integers, so it gives the exact answer with no rounding. If the modular output of your program matches the exact fraction reduced modulo the prime, the arithmetic is right. If it does not match, the bug is almost always in the model, not in the modular step.
This lesson covers the expected value as a weighted average, linearity of expectation, indicator variables, how to write a fraction as a number modulo p, and a DP over expectations. The C++ code is C++17, compiled with g++ -std=c++17 -O2 -Wall -Wextra using GCC 15.2.0. The Python snippets were run with Python 3.14.4. The output blocks are from our runs.
The expected value
The expected value of a random number is its average over the possible outcomes, weighted by how likely each outcome is. For a fair six-sided die, each face has probability 1/6, so
E[X] = 1 times 1/6 + 2 times 1/6 + 3 times 1/6 + 4 times 1/6 + 5 times 1/6 + 6 times 1/6 = 21/6 = 7/2
The expected value is 3.5, but no single roll ever shows 3.5. It is the long-run average over many rolls, not a value you should expect in one roll. Many wrong answers come from reading the expected value as the most likely result, which is a different thing.
Expected values are linear. Two facts follow from this definition, and they cover most of what contest problems need.
- E[a X + b] = a E[X] + b, for constants a and b.
- E[X + Y] = E[X] + E[Y], for any two random numbers X and Y, whether they depend on each other or not.
The second fact is the important one. It holds with no independence assumption, and most of the speed in these problems comes from it.
Linearity and indicator variables
An indicator variable is a random number that is 1 when an event happens and 0 when it does not. Its expected value is just the probability of the event, because 1 times P(event) + 0 times P(no event) = P(event).
Any count can be written as a sum of indicators. If X counts the events A_1, A_2, ..., A_n that happen, then X = I_1 + I_2 + ... + I_n, and linearity gives
E[X] = P(A_1) + P(A_2) + ... + P(A_n)
This is the whole method. You do not need the distribution of X. You only need the probability of each simple event.
A concrete example. Shuffle n cards and count the fixed points, meaning the cards that land on their own starting position. Let I_i be 1 if card i is a fixed point. The chance that card i stays put is 1/n, because the shuffle is uniform. So
E[number of fixed points] = n times 1/n = 1
The answer is 1 for every n, even though the events "card 1 is fixed" and "card 2 is fixed" are not independent. Two cards cannot both be fixed unless the rest of the shuffle is fixed too. Linearity does not care.
We checked this by listing every permutation for n from 1 to 7.
Output from our run:
n=1: permutations=1, average fixed points=1
n=2: permutations=2, average fixed points=1
n=3: permutations=6, average fixed points=1
n=4: permutations=24, average fixed points=1
n=5: permutations=120, average fixed points=1
n=6: permutations=720, average fixed points=1
n=7: permutations=5040, average fixed points=1The brute force agrees with the formula for all seven sizes. Two dependent quantities, here the fixed points of different cards, still add.
A second example shows the trap of multiplying expected values. Roll one die to get X. Let Y be the same roll, so Y = X. Then E[X] = 3.5 and E[Y] = 3.5, and E[X + Y] = 7. That is linearity, and it is right. But E[X times Y] = E[X squared] = 91/6, which is not 3.5 times 3.5 = 12.25. The product rule E[XY] = E[X] E[Y] needs the two parts to be independent, and here they are not. Use linearity for sums. Only use the product rule when you can prove independence.
Fractions modulo a prime
Contest answers are fractions. You want the numerator times the inverse of the denominator, taken modulo a prime p. The inverse of b modulo p is a number c with b times c = 1 modulo p. For a prime p, the inverse is b raised to the power p - 2, by Fermat's little theorem. The inverse exists exactly when p does not divide b. So a fraction a/b is written as a times b^(p-2) modulo p, and this needs b not divisible by p.
Take 7/2 modulo 998244353. The inverse of 2 is (998244353 + 1) / 2 = 499122177. Then 7 times 499122177 modulo 998244353 is 499122180. To check, 2 times 499122180 = 998244360, which is 7 more than a multiple of 998244353, so the result is 7 modulo p. The printed answer is 499122180, not 3.5 and not 7/2.
Two rules keep this safe.
- Never print a decimal number when the statement asks for a fraction modulo p. A decimal is a different number.
- Put each division through the inverse. Do not divide an integer by another integer inside a formula that is only exact over the rationals.
The inverse of many numbers at once can be done in a line each, with the formula inv[i] = -(p / i) times inv[p mod i] modulo p, which works because p mod i is smaller than i. The program below uses it.
DP over expectations
Many expected-value problems are a process with states. From each state you move to another state with some probability, or you stay where you are. Let E[s] be the expected number of further steps from state s. Then E[s] is a weighted average of 1 plus the E values of the states you can reach. These equations look circular when a state can move to itself. The fix is to solve the equation for E[s] by hand, which is what a DP over expectations means in practice.
A worked example, with a loyalty card. A shop has n different stamps, and each visit gives one stamp chosen uniformly at random from the n kinds. Lena already has m different stamps. How many visits does she need on average to collect all n? Call k the number of different stamps she has. Let E[k] be the expected number of visits still needed. With probability k/n the visit gives a stamp she already has, and she stays in state k. With probability (n - k)/n the visit gives a new stamp, and she moves to state k + 1. So
E[k] = 1 + (k/n) E[k] + ((n - k)/n) E[k + 1]
The equation has E[k] on both sides. Move the term with E[k] to the left:
E[k] (1 - k/n) = 1 + ((n - k)/n) E[k + 1]
Since 1 - k/n = (n - k)/n, multiply both sides by n/(n - k):
E[k] = n/(n - k) + E[k + 1], with E[n] = 0
Now the answer is a sum with no circularity. Starting from k = m,
E[m] = n/(n - m) + n/(n - m - 1) + ... + n/1
For n = 3 and m = 0 this is 3/3 + 3/2 + 3/1 = 11/2. That is the classic coupon collector result, and the recurrence reproduces it. We did not derive a new formula here. The point of the example is the method: solve the self-loop by hand, then the rest is a sum.
Here is the exact-fraction version of that sum, which is the check we use in the code section.
Output from our run:
11/2 9/2 25/3The code
The C++ program below reads n and m, and prints the expected number of visits as a fraction modulo 998244353. It computes the inverses of 1 to n with the linear recurrence, then adds the terms.
#include <bits/stdc++.h>
using namespace std;
using u64 = unsigned long long;
const u64 MOD = 998244353;
// Answer: expected number of visits until all n stamps are collected,
// starting with m distinct stamps. The answer is n * sum over k = m..n-1 of 1 / (n - k),
// which is the same as sum over k of n / (n - k).
int main() {
long long n, m;
cin >> n >> m;
// Inverses of 1..n by the linear recurrence inv[i] = -(MOD / i) * inv[MOD % i].
vector<u64> inv(n + 1, 1);
for (long long i = 2; i <= n; i++)
inv[i] = (MOD - (MOD / i) * inv[MOD % i] % MOD) % MOD;
u64 ans = 0;
for (long long k = m; k <= n - 1; k++)
ans = (ans + (u64)n % MOD * inv[n - k]) % MOD;
cout << ans << '\n';
return 0;
}Output from our run:
$ printf '3 0\n' | ./expect_stamps
499122182
$ printf '3 1\n' | ./expect_stamps
499122181
$ printf '4 0\n' | ./expect_stamps
332748126
$ printf '1000000 0\n' | ./expect_stamps
697618461The first line is 11/2 modulo p. The second is 9/2 modulo p, and the third is 25/3 modulo p. You can check each one by multiplying back by the denominator. For example, 3 times 332748126 is 998244378, which is 25 more than p, so 3 times the answer is 25 modulo p.
The check against exact fractions is a separate script. It runs the program for every n from 1 to 30 and every m below n, which is 465 pairs, and compares the output with the exact fraction reduced modulo p.
from fractions import Fraction
import subprocess
MOD = 998244353
def exact(n, m):
E = Fraction(0)
for k in range(n - 1, m - 1, -1):
E = Fraction(n, n - k) + E
return E
def to_mod(f):
return f.numerator % MOD * pow(f.denominator % MOD, MOD - 2, MOD) % MOD
bad = 0
cases = 0
for n in range(1, 31):
for m in range(0, n):
out = subprocess.run(["./expect_stamps"], input=f"{n} {m}\n", capture_output=True, text=True, timeout=10).stdout.strip()
if int(out) != to_mod(exact(n, m)):
bad += 1
cases += 1
print("exact-fraction check:", cases, "cases,", bad, "mismatches")Output from our run:
exact-fraction check: 465 cases, 0 mismatchesWorked problem: a stamp card with a head start
Problem. A shop has n different stamps. Each visit gives one stamp, chosen uniformly at random from the n kinds, independent of earlier visits. Lena already has m different stamps, with 0 at most m at most n - 1. Print the expected number of visits until she has all n kinds, as a fraction modulo 998244353. The input is n and m on one line. Limits: n at most 10^6.
Idea. The state is the number of different stamps Lena has. Each visit either adds a new kind or repeats one. The recurrence E[k] = n/(n - k) + E[k + 1] removes the self-loop, and the answer is a sum from k = m to n - 1. We compute each 1/(n - k) with the inverse table.
Full solution. The main function in the code above does this in three lines of work. It builds the inverses of 1 to n, then adds n times inv[n - k] for each k from m to n - 1. The total cost is O(n) time and memory.
Test. We checked the program on every pair with n from 1 to 30 and m below n, 465 pairs in all, against the exact fraction computed with Python's Fraction. There were 0 mismatches. We also ran n = 3 and m = 0, which gives 11/2 exactly and 499122182 modulo p, and a large case n = 10^6, m = 0, which printed 697618461. We did not check the large case against an exact fraction, because the exact sum has about 10^6 terms and a huge denominator. Its correctness rests on the small checks and on the recurrence.
Complexity
Time O(n) for the inverse table and the sum, for n up to 10^6. Memory O(n). The modular inverse of each term is found by the linear recurrence, so there is no fast exponentiation per term.
Common mistakes
Printing a decimal. Symptom: the output is 3.5 or 0.333333 where the statement asks for a fraction modulo p. Convert the fraction with the modular inverse and print the integer.
Multiplying expected values of dependent parts. Symptom: the answer for E[XY] is wrong when X and Y come from the same random event. Use linearity for sums. Use E[X] E[Y] only when the two are independent, and say so in your notes.
Forgetting the self-loop. Symptom: the code loops forever, or the recurrence reads E[k] before it has been set. Move the self-loop term to the left side and divide, as the worked example does.
Dividing by a multiple of p. Symptom: the modular inverse is 0, and every later term is 0. The inverse of b exists only if p does not divide b. For the stamp problem, n - k is between 1 and n, so this cannot happen when n is below p.
Reading the expected value as the most likely outcome. Symptom: a tool answers "the most likely number of visits" with the expected value. These are different. The expected value is an average over many runs, and it need not be a value that ever occurs.
Trusting a float on small inputs. Symptom: the float answer looks right to six digits and the judge still rejects it. A fraction with a large denominator cannot be recovered from a float, so keep integers throughout.
Checklist
- Write the expected value as a sum over outcomes, or as a sum of indicator variables.
- Use linearity for sums, with no independence assumption.
- Use the product rule only when you can say why the two parts are independent.
- Write each probability as a fraction, then convert with the modular inverse.
- Check that every denominator is below p and not divisible by p.
- For a DP with a self-loop, solve the equation for the state by hand before coding.
- Compare the modular output with the exact fraction from Python's
fractionson small inputs, and check that the two agree after the modular reduction. - Test the boundary cases: the start state at the end, and the smallest n.
Exercise
A fair die with n faces is rolled until some face has come up for the second time, in any rolls, not necessarily in a row. Find the expected number of rolls, as a fraction modulo 998244353. Input: n, with 2 at most n at most 10^6. Write the solution with a DP over the number of distinct faces seen so far. Then check it against exact fractions for every n from 2 to 20, using Python's fractions. Print the number of cases and the number of mismatches. Keep the exact check as a separate function from the modular code.
Prompt to use with an AI tool, if you want help:
Write C++17 code that reads n, the number of faces of a fair die, and prints the
expected number of rolls until some face shows for the second time, as a fraction
modulo 998244353. Use a DP over d, the number of distinct faces seen so far:
F[n] = 1, and F[d] = 1 + ((n - d) / n) times F[d + 1]. The answer is 1 + F[1].
Write the division by n with the modular inverse. Also write a Python script
that computes the same value with fractions.Fraction for n from 2 to 20 and
compares the two. Print the number of cases and the mismatches.A good answer uses the modular inverse for every division by n, starts the DP at F[n] = 1 and not at 0, and keeps the exact check independent of the C++ code. For n = 2, the exact answer is 5/2, which is 499122179 modulo 998244353. If the tool prints 3 for n = 2, it has solved the "two in a row" version instead, which gives n + 1 = 3, so check the stop rule against the statement.
Show the answer
The expected number of rolls is 1 + F[1], with F[n] = 1 and F[d] = 1 + ((n - d)/n) F[d + 1]. The values are 5/2 for n = 2, 26/9 for n = 3, 103/32 for n = 4, 2194/625 for n = 5 and 1223/324 for n = 6. We checked the recurrence against the tail-sum formula, the sum over k of the probability that the first k rolls are all different, for n from 2 to 6, and both agree.
Practice problems
Problem 1. A fair coin is tossed until heads appears. What is the expected number of tosses?
Hint: Let E be the expected number of tosses. The first toss is heads with probability 1/2, otherwise you start again. This gives an equation with E on both sides.
Show the answer
- Write E = 1 + (1/2) E, so E = 2. This is the sum of 1/2 + 2/4 + 3/8 + ... over all tosses, which is 2.
Problem 2. The number of heads in 10 fair coin tosses. What is its expected value, and what is the expected value of its square?
Hint: The expected value is a sum of indicators. The square needs the variance, which for a sum of independent indicators is the sum of their variances.
Show the answer
The expected number of heads is 5. Each toss is an indicator with variance 1/4, and the tosses are independent, so the variance of the count is 10/4 = 5/2. The expected square is the variance plus the square of the mean, which is 5/2 + 25 = 55/2.
Problem 3. Compute 5/3 modulo 998244353.
Hint: The inverse of 3 modulo p is 332748118, because 3 times 332748118 is 998244354. Multiply by 5 and reduce.
Show the answer
Multiply 5 by 332748118 to get 1663740590, then reduce: 1663740590 minus 998244353 is 665496237. Check: 3 times 665496237 is 1996488711, which is 5 more than 2 times p. So the answer is 665496237.
Problem 4. Two cards are drawn from a shuffled deck of n cards without replacement. Let X be the number of cards among the two that are red, where the deck has r red cards. Find E[X] using indicators, and check the answer for n = 4 and r = 2.
Hint: Let I_1 and I_2 be the indicators for the first and second card being red. The chance that the first card is red is r/n, and the chance that the second card is red is also r/n, even without independence.
Show the answer
E[X] = 2r/n. For n = 4 and r = 2 this is 1. Check by listing: the deck has 2 red and 2 black cards, so there are 6 pairs. One pair has 2 red cards, four pairs have 1, and one pair has 0. The total is 2 + 4 = 6, and 6 / 6 = 1, which matches.
Summary
The expected value is a weighted average, and it can be found without knowing the full distribution. Linearity of expectation lets you add the expected values of simple indicator events, with no independence needed, and a DP over expectations becomes a sum once you solve the self-loop by hand. Answers are fractions, so write each one with a modular inverse, and check your program on small inputs against exact fractions. A float is not the answer the judge wants.