Probability of a number N being reduced to a nonpositive value after K trials

Here's a program to calculate the required probability

N = 2
M = 1
K = 3

count = [0] * (N+1)
prev = [0] * (N+1)

count[0] = 1 # empty set

for i in range(K):
    # move count to prev
    for index in range(N+1):
        prev[index] = count[index]
        count[index] = 0
    
    # calculate new counts
    for prevSum in range(N+1):
        for value in range(M+1):
            newSum = min(N, prevSum+value)
            count[newSum] += prev[prevSum]
            
ans = (count[N] / pow(M+1, K))
print(ans)

working code link

• Here we keep track of the count of number of sets that add upto a given sum in the count[] array
• Any set that adds upto a value greater than N is added to count[N]

How does this work?

• Initially, count[0] = 1, since we only have empty set {}
• (K = 1): try adding one element to all the existing sets: {} + 0, {} + 1. we get {0}, {1} so count[0] = 1, count[1] = 1
• (K = 2): Now again add one element to all the existing sets: {0} + 0, {0} + 1, {1} + 0, {1} + 1. we get {0,0},{0,1},{1,0},{1,1} so count[0] = 1, count[1] = 2, count[2] = 1
• (K = 3): Now again add one element to all the existing sets. We would get {0,0,0}, {0,0,1}, {0,1,0}, {1,0,0}, {0, 1, 1}, {1, 0, 1}, {1,1,0}, {1,1,1}. so count[0] = 1, count[1] = 3, count[2] = 4
• here in the last step, we also add {1,1,1} to count[2] because we add all sets whose sum is >= N to count[N]
• Finally, to compute the probability we divide count[N](count of all sets whose sum is >=N) with the count of all possible sets i.e, (M+1)^K

The complexity is O(N*M*K). In worst case N=M*K, so the time complexity can be rewritten as: O((M*K)^2)

Optimisation 1:

If you write down the count[] array for each iteration of K you can find an interesting observation:

M=1

sum: 0 1 2 3 4
K=0: 1 0 0 0 0 (if empty consider value as 0 from now on)
K=1: 1 1
K=2: 1 2 1
K=3: 1 3 3 1
K=4: 1 4 6 4 1


M=2

sum: 0  1  2  3  4  5  6  7  8
K=0: 1
K=1: 1  1  1
K=2: 1  2  3  2  1
K=3: 1  3  6  7  6  3  1
K=4: 1  4 10 16 19 16 10  4  1

the observation here is:

By maintaining a rolling sum of previous M values, we can write optimised version of the code:

N = 2
M = 1
K = 3
maxValue = M*K

count = [0] * (maxValue+1)
prev = [0] * (maxValue+1)

count[0] = 1 # empty set

for i in range(K):
    # move count to prev
    for index in range(maxValue+1):
        prev[index] = count[index]
        count[index] = 0
    
    rollingSum = 0
    
    # calculate new counts
    for Sum in range(maxValue+1):
        rollingSum += prev[Sum]
        if (Sum > M):
            rollingSum -= prev[Sum - (M + 1)]
        count[Sum] = rollingSum
            

# add all counts of sets whose sum is >= N
ans = sum(count[N:]) / pow(M+1,K)
print(ans)

working code link

The time complexity of this approach is O(M*(K^2))

I will focus at the algorithm level only here.

The problem is equivalent to calculate the probability than after K steps, the sum of the values is greater than N.

Let ut call p the probability that at a given step, on element is selected

p = 1/(M+1)

The idea is to represent the probability set of one trial in a polynomial form:

A[x] = p * (1 + x^2 + x^3 + ... + x^{M-1})

Then after K steps, the probability set is given by:

F[x] = p^K * (1 + x^2 + x^3 + ... + x^{M-1})^K = p^K A[x]^K

At the end, if F[x] = sum_i p_i x^i, then the probability is given by:

P = sum_{i >= N} p_i

The problem is then to calculate the polynomial power A[x]^k.

A first attempt consits in calculating it recursively: then we obtained something equivalent to be solution provided in the first answer, same complexity.

A second attempt consist in implicitely considering the binary representation of K.

Pseudo-code:

F[x] = 1
T[x] = A[x]
Power = K
while (Power != 0)
    if (Power mod 2 == 1) F[x] = F[x] * T[x]
    T[x] = T[x] * T[x]
    Power = Power / 2
end while

The number of steps is equel to log2(K).
At each iteration, the degree of the polynomial is multiplied by two: the complexity is dominated by the last step, where the order of the polynomial is equal to KM.
This last step has a complexity O(K^2 M^2).

If C = K^2 M^2, The global complexity of the power calculation is O(C + C/4 + C/8 + ...) = O(C) = O(K^2M^2).

This method doesn't seem to provide any advantage with respect to the first one (I may have made an error in the complexity evaluation!).

A third method consists in considering that a polynomial multiplication is equivalent to a convolution, and can be performed via a FFT process. As the final size is O(KM), then the complexity is O(M K (log M + log K)).

Describing this method in details here would take much time. You could find many references on internet on this subject, for example here

Side note: a pity not being able to insert latex math formula here ...