chbp.py

f=open("rosalind_chbp.txt","rU")
taxa=f.readline().split()
taxon_names=dict(zip(range(len(taxa)),taxa))
characters=[x.strip() for x in f.readlines()]
def LR(leafs,ch):
    left=[leafs[i] for i in range(len(leafs)) if ch[leafs[i]]==ch[leafs[0]]]
    right=[leafs[i] for i in range(len(leafs)) if ch[leafs[i]]!=ch[leafs[0]]]
    return left,right


def build_subtree(leafs,characters):
    if len(characters)==0:
        return ','.join([taxon_names[leafs[i]] for i in range(len(leafs))])
    if len(leafs)==2:
        return '(%s,%s)'%(taxon_names[leafs[0]],taxon_names[leafs[1]])
    elif len(leafs)==1:
        return '%s'%taxon_names[leafs[0]]
    elif len(leafs)==0:
        return ''
    (left,right)=LR(leafs,characters[0])
    print 'Left',[taxon_names[x] for x in left]
    print 'Right',[taxon_names[x] for x in right]
    left_tree=build_subtree(left,characters[1:])
    right_tree=build_subtree(right,characters[1:])
    if len(left_tree)==0:
        return '%s'%right_tree
    elif len(right_tree)==0:
        return '%s'%left_tree
    else:
        return '(%s,%s)'%(left_tree,right_tree)

def build_tree(leafs,characters):
    (left,right)=LR(leafs,characters[0])
    left_tree=build_subtree(left,characters[1:])
    right_tree=build_subtree(right,characters[1:])
    if left_tree[0]=='(':
        left_tree=left_tree[1:-1]
    elif right_tree[0]=='(':
        right_tree=right_tree[1:-1]
    return '(%s,%s)'%(left_tree,right_tree)

s=build_tree(range(len(taxa)),characters)