Polya Urn Expected Red Count

Problem

In a Polya urn starting with red red and blue blue balls, each draw adds one ball of the drawn color. The expected red fraction stays constant, so E[red after n steps] = red + n * red / (red + blue).

Given the initial counts and n, return the expected number of red balls.

Hint

The expected increment at any step is the initial red fraction.

Test cases

case 1

in: [1,1,2]

expected: 2

case 2

in: [3,1,4]

expected: 6

case 3

in: [2,2,0]

expected: 2

Starter codePython

def polya_urn_expected_red(red, blue, n):
    # Your code here
    pass

3 test cases · run your code against them in the editor