mend.py

from nwck_class_rooted import Newick_Tree_Rooted

class Pedigree(Newick_Tree_Rooted):
    def __init__(self,s,flag=1):
        Newick_Tree_Rooted.__init__(self,s)
        self.flag=flag

popdict={'AA':[1,0,0],'Aa':[0,1,0],'aa':[0,0,1]}
pops={}


with open("rosalind_mend.txt","rU") as f:
    s=f.read().strip()

T=Pedigree(s)

def population(T,node):
    new=[0,0,0]
    if len(T.dtree[node])==0:
        pops[node]=popdict[T.label_name[node]]
    else:
        population(T,T.dtree[node][0])
        population(T,T.dtree[node][1])
        left=pops[T.dtree[node][0]]
        right=pops[T.dtree[node][1]]
        new[0]=left[0]*right[0]+0.5*(left[0]*right[1]+right[0]*left[1])+0.25*left[1]*right[1]
        new[1]=0.5*(left[0]*right[1]+right[0]*left[1])+0.5*left[1]*right[1]+0.5*(left[1]*right[2]+right[1]*left[2])+(left[0]*right[2]+left[2]*right[0])
        new[2]=left[2]*right[2]+0.5*(left[1]*right[2]+right[1]*left[2])+0.25*left[1]*right[1]
        pops[node]=new
    
population(T,T.root)