probability puzzle, expected payoff for game

This was a question at a programming contest that finished recently at interviewstreet:

Alice and Bob play a game. The operations at round $i$ ($i ≥ 1$) is as follows:

• Alice pays Bob $2i - 1$ dollars,
• Alice tosses a biased coin,
• If the result of the coin was heads for $k$ consecutive rounds, the game stops, otherwise the game continues.

Given $k$ and the probablity $p$ that the outcome of a toss is heads, your program should find the expected number of dollars Alice pays Bob, and also the expected number of rounds played.

I had gotten the first part of the problem, i.e the expected number of rounds of the game. I got it with the following code

if __name__ == '__main__':
t = int(raw_input())
while t :
t -= 1
temp = str(raw_input())
p,k = temp.split(' ')
p = float(p)
k = int(k)
num = k * (p**k)
den = 1
q = 1.0 - p
for N in range(1,k+1):
den = den - ((p**(N-1))*q)
num = num + (N*(p**(N-1))*q)
#print (N*(q**N))

print int(num/den)


But the second part of the problem is still puzzling me, i.e the expected number of dollars Alice pays bob. How can expected payoff be calculated?

-

Call $N$ the number of rounds played, then the number of dollars paid by Alice is $N^2$. The usual one-step recursion yields $$\mathrm E_i(s^N)=s(p\mathrm E_{i+1}(s^N)+(1-p)\mathrm E_0(s^N)),$$ for every $0\leqslant i\leqslant k-1$, with $\mathrm E_k(s^N)=1$. One is after $\mathrm E_0(N)$ and $\mathrm E_0(N^2)$ and the linear system above yields $\mathrm E_0(s^N)=u(s)$ with $$u(s)=\frac{(1-ps)(ps)^k}{1-s+(1-p)s(ps)^k}=1-(1-s)\frac{1-(ps)^k}{1-s+(1-p)s(ps)^k},$$ hence $u'(1)=\mathrm E_0(N)$, that is, $$\mathrm E_0(N)=\frac{1-p^k}{(1-p)p^k}.$$ Similarly, $u''(1)=\mathrm E_0(N(N-1))$. This yields a slightly more complicated, but explicit, formula for $\mathrm E_0(N^2)$.

-
Thank you for your answer. I don't quite understand, some of the notation, what is s, and how did you calculate u(s)? –  rohanag May 9 '12 at 11:25
The argument $s$ can be any number such that $|s|\leqslant1$. One computes $u(s)$ by considering, for each such $s$, the affine system the vector $(\mathrm E_i(s^N))_{0\leqslant i\leqslant k}$ solves. –  Did May 9 '12 at 13:23