Bug in distance_between function

Hi,

I noticed that the distance between function would return the "u and v are not in the same Tree" error if v is an ancestor of u. I was able to fix it by inserting one line of code. Here is the updated version of the function:

def distance_between(u, v):
        if u == v:
            return 0.
        elif u == v.parent:
            return v.edge_length
        elif v == u.parent:
            return u.edge_length
        u_dists = {u:0.}; v_dists = {v:0.}
        c = u; p = u.parent # u traversal
        while p is not None:
            u_dists[p] = u_dists[c]
            if c.edge_length is not None:
                u_dists[p] += c.edge_length
            if p == v:
                return u_dists[p]
            c = p; p = p.parent
        c = v; p = v.parent # v traversal
        while p is not None:
            v_dists[p] = v_dists[c]
            if c.edge_length is not None:
                v_dists[p] += c.edge_length
            if p in u_dists:
                return u_dists[p] + v_dists[p]
            c = p; p = p.parent
        raise RuntimeError("u and v are not in the same Tree")

The only change that I have made here is adding these two lines of code: if p == v: return u_dists[p] (lines 14 and 15). Hope this helps.

Best, Shayesteh

niemasd / TreeSwift

Bug in distance_between function #34