Optimal Stopping for Die Rolls That Bust on Perfect Squares
Company: Jane Street
Role: Data Scientist
Category: Coding & Algorithms
Difficulty: medium
Interview Round: Onsite
You repeatedly roll a fair six-sided die and keep a running cumulative sum $S$ of all the values rolled so far (starting from $S = 0$ before the first roll).
- After any roll, you may choose to **stop** and walk away with a payoff equal to the current sum $S$.
- However, if at any point the cumulative sum $S$ lands exactly on a positive perfect square ($1, 4, 9, 16, 25, \dots$), you immediately **lose everything**: the game ends with payoff $0$.
- If you keep rolling and have not busted, the game continues indefinitely — there is no limit on the number of rolls.
Determine the optimal strategy: for which values of the running sum should you keep rolling, and at which values should you stop and take the money? Justify why the strategy is optimal, and compute the expected payoff of the game under optimal play, starting from a sum of $0$.
Quick Answer: This question evaluates probabilistic reasoning and optimal stopping competency, testing understanding of Markov decision processes, expected-value calculations, and risk trade-offs in stochastic processes.
You repeatedly roll a fair `sides`-faced die (each face 1..`sides` equally likely) and keep a running cumulative sum `S`, starting at `S = 0` before the first roll.
- After any roll you may **stop** and take the current sum `S` as your payoff.
- If the running sum ever lands *exactly* on a **positive perfect square** (1, 4, 9, 16, 25, ...) you **bust**: the game ends immediately with payoff 0. (The starting sum 0 is not a bust.)
- Otherwise the game may continue with no roll limit.
Because the true game is infinite-horizon, the challenge uses a truncation parameter `cap`: for any running sum `S > cap` you are *forced to stop* and take `S` (equivalently, the value of such a state is `S`). For a large enough `cap` this equals the true optimal value (for the standard 6-sided die the value already stabilises for `cap >= 100`).
Under optimal play, return the **expected payoff starting from `S = 0`** as a floating-point number (accepted within 1e-6).
Optimal-play recurrence (let `V(S)` be the optimal expected payoff at a decision point with sum `S`):
- `V(S) = S` for every `S > cap` (forced stop).
- For `0 <= S <= cap` that is **not** a positive perfect square:
`V(S) = max( S, (1/sides) * sum over d=1..sides of value(S + d) )`,
where `value(x) = 0` if `x` is a positive perfect square, else `value(x) = V(x)`.
Return `V(0)`.
Constraints
- 1 <= sides <= 10
- 1 <= cap <= 2000
- The die is fair: each face in 1..sides has probability 1/sides.
- Only positive perfect squares (1, 4, 9, 16, ...) bust; the starting sum 0 does not.
- Return a float; answers within 1e-6 of the reference are accepted.
Examples
Input: (6, 300)
Expected Output: 7.175499312767105
Explanation: Canonical game: fair 6-sided die, cap large enough that the value equals the true infinite-horizon optimum (~7.1755).
Input: (6, 50)
Expected Output: 7.158761036815606
Explanation: Same die but a smaller truncation cap (300 -> 50) yields a slightly different value, exercising the S>cap boundary.
Hints
- Solve backwards. There is no upper bound in the real game, so use the given `cap`: every state with sum S > cap has value exactly S (forced stop).
- For a decision state S <= cap that is not a square, V(S) = max(S, average of value(S+d) over d = 1..sides), where a landing on a positive perfect square contributes 0.
- Because stopping is always available, V(S) >= S. Whenever the next `sides` sums contain no perfect square, rolling strictly beats stopping (each successor is already >= its own index), so the only real decisions happen just below a square.
- Use integer square-root (e.g. math.isqrt) to test whether x is a perfect square in O(1); iterate S from cap down to 0 filling a V array.