Problem statement

https://leetcode.com/problems/toss-strange-coins/

Solution

Let dp(i, j) be probability to have j heads if we already used coin upto number i inclusive. Then we have border cases:

  1. If j < 0: we can not have negative number of head or i - j < -1, that is we have more coins that we can get heads.
  2. If i == -1, we return 1, it means we do not have coins and number of heads we need is 0.
  3. Finally to evaluate dp(i, j) we need to look at dp(i-1, j) and dp(i-1, j-1).

Complexity

It is O(n*k) for time and space. Notice, that we need to be careful with lru_cache, sometimes it will give TLE, if we put None as parameter, we can try to use 10000 for example and it will pass.

Code

class Solution:
    def probabilityOfHeads(self, P, target):
        @lru_cache(None)
        
        def dp(i, j):
            if j < 0 or i - j < -1: return 0
            if i == -1: return 1
            return dp(i-1, j)*(1-P[i]) + dp(i-1, j-1)*P[i]
        
        return dp(len(P) - 1, target)