Solve LCA taks with binary path method

LCA.py

## Solve LCA taks with binary path method
## Complexity is <O(NlogN), O(logN)>

tree = {
    0:  [1,2],
    1:  [3,4,5],
    2:  [6,7,8,9],
    3:  [10,11,12,13],
    4:  [14,15],
    5:  [16,17,18,19,20],
    6:  [21,22,23,24],
    7:  [25],
    8:  [26],
    9:  [27],
    10: [28]
}

## Preprocessing
tin, tout = [0]*29, [0]*29
P = [[None]*5 for _ in xrange(29)]

## Function to check if L node is ancestor for R
def ancestor(l,r):
    return tin[l] < tin[r] and tout[l] > tout[r]

## Tree traversal
def dfs(node, p, step):
    tin[node] = step

    P[node][0], i = p, 1
    while P[node][i-1] is not None:
        P[node][i] = P[P[node][i-1]][i-1]
        i += 1

    for t in tree.get(node, ()):
        step = dfs(t, node, step+1)

    tout[node] = step+1
    return step+1

dfs(0, None, 0)

## Check that ancestor test works find
assert ancestor(1,2) == False
assert ancestor(0,1) == True
assert ancestor(1,0) == False
assert ancestor(0, 28) == True
assert ancestor(1, 15) == True

## LCA calculation
def lca(l, r):
    if ancestor(l,r): return l
    if ancestor(r,l): return r

    p = l
    for i in range(4,-1,-1):
        if P[p][i] is not None and not ancestor(P[p][i], r):
            p = P[p][i]
        i -= 1

    return P[p][0] or 0

## Few tests to check LCA correctness
assert lca(4, 5) == 1
assert lca(1, 2) == 0
assert lca(28, 10) == 10
assert lca(28, 2) == 0