Calculate Probabilites of victory in a coin flip game with n flips

Question

In this game, a coin is flipped n times. Player 1 gains a point when two heads are flipped in a row, Player 2 gains a point when a head is flipped and followed by a tails.

For example, in the string "HTTTHHHT," the outcome would be a tie because Player 1 received two points from the "HHH" substring and Player 2 also received two points from the "HT...HT" substrings.

I have developed a brute-force algorithm to calculate the probability of each player winning or a tie, but it cannot handle large values. Since there are 2^n permutations of an n length string, my program will fail when 2^n exceeds the integer limit. Even if I changed it to be able to properly handle larger values, it would take incredibly long to calculate larger n values.

How would I be able to develop a better algorithm as opposed to generating each permutation and determining who wins?

Here is my current algorithm, assume heads are 0 and tails are 1 (Alice: HH, Bob: HT):

public class Probability {
    public static void main(String[] args) {
        int n = 5;
        int[] winners = Permutations(n);
        float pA = winners[0] / (float) (1 << n);
        float pB = winners[1] / (float) (1 << n);
        float pT = winners[2] / (float) (1 << n);

        System.out.printf("%d Tosses. P(Alice): %.3f. P(Bob): %.3f. P(Tie): %.3f\n", n, pA, pB, pT);
    }

    public static int[] Permutations(int n) {
        int amtPerm = 1 << n; // How many permutations to generate
        int[] perm = new int[n]; // Fills with all 0 by default
        int[] winners = {1, 0, 0}; // Alice, Bob, Tie (Counted the first permutation of all 0s as an Alice win)

        for(int i = 1; i < amtPerm; i++) {
            int j = n-1;
            while(true) {
                if(perm[j] == 0) {
                    perm[j++] = 1;
                    for(; j < n; j++) { // set all values after j to 0
                        perm[j] = 0;
                    }
                    break;
                } else {
                    j--;
                }
            }
            // Traverse the new permutation to determine a winner
            int alice = 0, bob = 0;
            for(int k = 1; k < n; k++) {
                if(perm[k-1] == 0 && perm[k] == 0) alice++;
                else if(perm[k-1] == 0 && perm[k] == 1) bob++;
            }
            if(alice > bob) winners[0] += 1;
            else if(bob > alice) winners[1] += 1;
            else winners[2] += 1;
        }

        return winners;
    }
}

As mentioned above, the algorithm works as expected, but I would like to expand this to calculate probabilities with larger values of n. If I input 1000 as my value of n, I get this as my probability results: "P(Alice): 1.000, P(Bob): 0.000, P(Tie): 0.000"

Answer 1

The tricky part, as I'm sure you've noticed, that you can't just break the problem into smaller subproblems counting wins/losses, you actually need the point breakdowns from those subproblems.

You could absolutely track all assignments of Alice's and Bob's points for each individual subproblem size. In fact, you could do better, by cutting each problem into two almost-equal halves (+/-1), you can bound the number of unique subproblem sizes in O(log N). For n = 2*k + 1, recurse both halves on k, and for one half just increment the counts accordingly.

For this approach, I'd recommend a recursive formula, taking five parameters (two of which are just binary choices): # of flips, first flip, last flip, bob points, alice points- returning the associated count. With proper caching (or building bottom->top in DP fashion), this would take O(n^2 log n) time and space.

We can do better however. Instead, consider a different perspective, where we frame the problem as partition problem. For our n, we'll consider each value b, representing the number of points bob scores (specifically, times HT appears). For simplicity, also assume the flipping starts with heads (we can adjust the calculation for tails later). If we partition the flips into consecutive groups of identical flips, bob's points will be half (floored) the number of partitions.

We could easily count the number of such partitions, as its just the stars and bars problem, but we need to be able to also specify alice's points a. For a specific value of a and the number of partitions P, we can calculate the exact number of H and T occurrences. We can then count the number of ways partition each of these two into the desired number of (non-empty) partitions (ceil(P/2), floor(P/2)) via n-choose-k as with the stars and bars problem. The product of these counts is the total number of ways to partition the flips with a and b points each.

This will take O(N^2) time overall (assuming you cache / build a table for the n choose k calculation), ignoring integer multiplication- if you end up working with large enough N that that matters, then you'll need to factor in a O(log(N) log(log(N)) log(log(log(N))) for Schonhage-Strassen multiplication.

Lastly, if you want a sub-quadratic time algorithm, I believe this is possible, but its a lot of math and a lot of work. You can represent the difference between alice / bob's points in the exponent of a polynomial, c*x^(a-b), where c is the count of such point assignments for our input n, and x is an unbound variable. We can then write two expressions, relating the polynomials for a given n to those for n+1 for the cases where n ended with heads H(x) or where it ended with tails T(x):

H(n+1, x) = H(n, x) * x + T(n, x)
T(n+1, x) = H(n, x) / x + T(n, x)

This can be rewritten as a transformation matrix

    |  x     1  | 
M = |           |
    | 1/x    1  | 

We can raise M to a power (in our case, n) to apply it multiple times, and then multiply it with the starting vector [x^0, x^0] to get a pair polynomials representing our total counts for each point difference. We can raise this matrix to the nth power in O(log n) steps using exponentiation by squaring. We can multiply the polynomials in each matrix in O(n log n) time using the fast fourier transform (FFT), we'll just scale them up by a factor of x^n before hand and scale them down later to keep the exponents non-negative.

This means we can compute our polynomials in O(n log^2 n) time, or in reality O(n log n) time, since the last multiplication dominates the work. Then its just a matter of adding the two polynomials, dividing out that factor of x^n, and then adding up coefficients. The sum of the coefficients of terms with exponent greater than zero will be the count of times alice wins. Those less than zero will be the times bob wins, and the constant term will be the draws. This makes for an O(n log n) time approach overall, with a large constant of course.

---

Side note: I haven't had time to code any of these solutions up yet to validate their correctness nor gather performance benchmarks. If you see any issues or have any concerns, please call them out and I'll address them as soon as possible.

---

Update

Here's an implementation of the FFT-driven solution which makes use of numpy for much of the math implementation:

import numpy as np
from numpy.linalg import matrix_power
from numpy.polynomial.polynomial import Polynomial

def solve(n):
    matrix = np.array([[Polynomial([0, 0, 1]), Polynomial([0, 1])],
                       [Polynomial([1, 0, 0]), Polynomial([0, 1])]])
    coeffs = matrix_power(matrix, n - 1).sum().coef
    coeff_parts = np.split(coeffs, [n - 1, n])
    
    bob, draw, alice = [round(x.sum()) for x in coeff_parts]
    return alice, draw, bob

Example:

>>> solve(19)
(210934, 68568, 244786)

I've checked this against the results of a simple brute force solver, so its is in fact correct, and assuming that numpy is using the FFT to do polynomial multiplication under the hood, its should be O(n log n), and it is indeed pretty fast in practice. For n=50, it takes 0.003 sec on my machine.

Only issue is that the results get so large so quickly that the floating point arithmetic used in the FFT lacks the precision required to convert it back to integers. You could switch to the number theoretic transform (NTT) to avoid floating point, replacing numpy's Polynomial with your own. You'll also introduce a O(K*log(K)*log(log(K))) factor for Schonhage-Strassen integer multiplication, where K is the number of bits in the output integers.

---

Update 2

Here is code for the other solution involving counting partitions:

from math import comb

def nck(n, k):
    if n < 0 or k < 0:
        return 0
    return comb(n, k)

def calc(i, j, n):
    return sum(nck(ph + j - 1, ph - 1) * nck(n - ph - j - 1, pt - 1)
        for ph in range(i, i + 2) for pt in range(i, i + 2))

def solve_alt(n):
    draw = 1 + sum(calc(i, i, n)
            for i in range(n // 3 + 3))
    bob = sum(calc(i, j, n)
            for i in range(n // 2 + 2)
            for j in range(min(i, n - 2 * i + 1)))
    alice = 2 ** n - draw - bob
    return alice, draw, bob

Again, you'll probably want to use your own nck / comb which caches intermediate computations to result in less multiplications for a faster solver. This solution is already pretty fast, with n=50 taking 0.001 sec on my machine, but it will scale worse than the FFT-driven approach above, however it will not experience any accuracy issues unlike the FFT approach since this only deals in integers.