wfmd.py

import numpy as np
from scipy.misc import comb

with open("rosalind_wfmd.txt","rU") as f:
    line=f.readline().strip()

nums=line.split()
flts=map(float,nums)
    
popsize,dominants,generations,k=flts[0],flts[1],flts[2],flts[3]
N=2*popsize
print N,dominants,generations,k

m=np.zeros((N+1,N+1))
for i in np.arange(N+1):
    for j in np.arange(N+1):
        m[j,i]=comb(N,j)*(i/N)**j*((N-i)/N)**(N-j)

pop=np.zeros(N+1)
pop[dominants]=1
newpop=pop
for i in np.arange(generations):
    newpop=m.dot(newpop)
    print sum(newpop)
    

print round(sum(newpop[0:N-k+1]),3)