Challenge: Bear and Dancing

Subdomeniu: Probability (probability)

Scor cont: 100.0 / 100

Submission status: Accepted

Submission score: 1.0

Submission ID: 464778671

Limbaj: python3

Link challenge: https://www.hackerrank.com/challenges/bear-and-dancing/problem

Cerinta

Bear Limak is a dance teacher.
Today is the first day of the course.
The course will take one or more days.
Your task will be to calculate the expected value of the number of dances in the course.

There are $n$ boys and $m$ girls.
A classroom is very small and thus only one pair can dance at each moment.
For each new dance Limak chooses uniformly at random one boy and one girl.
The chosen pair will dance, unless the following will happen.

It's possible that the chosen pair has already danced with each other on the same day.
Then, with probability $r$ they will now get upset about it and they will refuse to dance (but otherwise they dance like a normal pair).
In such a situation Limak will apologize them and there will be no more dances on that day.
Classes will start again on the next day though, and Limak won't care who danced the day before and who got angry.
So, the situation will be exactly as on the first day.

Limaks waits for the possibility to say _"Nice, kids. Every person has danced today. The course is over!"_.
So, the course ends immediately when there is a situation that every person has danced on that day.
What is the expected value of the number of dances in the course?

Input Format

The only line of the input contains two integers $n$, $m$, and one real number $r$.

**Constraints**  

- $1 \le n \le 30$  
- $1 \le m \le 30$  
- $0.001 \le r \le 0.1$  
- $r$ is given with at most $6$ places after the decimal point.

Output Format

Find and print the expected value of the number of dances in the course.
The answer will be considered correct if the absolute or relative error doesn't exceed $10^{-6}$.

**Sample Input 0**

    1 2 0.062812
    
**Sample Output 0**
    
    3.0000000000

**Sample Input 1**

    2 3 0.075
    
**Sample Output 1**
    
    5.8901549035


**Sample Input 2**

    2 2 0.05
    
**Sample Output 2**
    
    3.6885245902

Cod sursa

import sys

def main():
    data = sys.stdin.read().strip().split()
    if not data:
        return
    n = int(data[0])
    m = int(data[1])
    r = float(data[2])

    NM = n * m

    nextFail = [[0.0] * (m + 2) for _ in range(n + 2)]
    nextExp = [[0.0] * (m + 2) for _ in range(n + 2)]

    for k in range(NM, -1, -1):
        pr = k / NM
        pup = pr * r
        den = 1.0 - pr * (1.0 - r)

        curFail = [[0.0] * (m + 2) for _ in range(n + 2)]
        curExp = [[0.0] * (m + 2) for _ in range(n + 2)]

        for a in range(n, -1, -1):
            na = n - a
            for b in range(m, -1, -1):
                if a == n and b == m:
                    continue

                if k > a * b:
                    continue

                mb = m - b

                p_new_new = (na * mb) / NM
                p_old_new = (a * mb) / NM
                p_new_old = (na * b) / NM
                p_old_old = (a * b - k) / NM

                numF = pup
                numE = 1.0 - pup

                if p_new_new:
                    numF += p_new_new * nextFail[a + 1][b + 1]
                    numE += p_new_new * nextExp[a + 1][b + 1]
                if p_old_new:
                    numF += p_old_new * nextFail[a][b + 1]
                    numE += p_old_new * nextExp[a][b + 1]
                if p_new_old:
                    numF += p_new_old * nextFail[a + 1][b]
                    numE += p_new_old * nextExp[a + 1][b]
                if p_old_old:
                    numF += p_old_old * nextFail[a][b]
                    numE += p_old_old * nextExp[a][b]

                curFail[a][b] = numF / den
                curExp[a][b] = numE / den

        nextFail, nextExp = curFail, curExp

    p_fail = nextFail[0][0]
    e_day = nextExp[0][0]
    p_succ = 1.0 - p_fail

    ans = e_day / p_succ
    print(f"{ans:.10f}")

if __name__ == "__main__":
    main()
HackerRank Probability – Bear and Dancing