获取N-ary树中较大祖先的数量

Question

We have a tree with nodes numbered from 1 to N . It is given to us in the format of 3 arrays vals , edge1 , edge2 . vals[i] has the value of node i + 1 . And there is an edge from edge1[i] to edge2[i]. The tree is rooted at 1.

I want to calculate the number of ancestors that are larger than the current nodes value for every node in the tree and return them in the an array.

This is a working algorithm:

def get_larger_ancestors(vals, edge1, edge2):
    tree = defaultdict(set)
    for u, v in zip(B, C):
        tree[u].add(v)
        tree[v].add(u)

    parents = {1:None}
    que = deque([1])
    seen = set()
    while que:
        current = que.popleft()
        seen.add(current)
        for node in tree[current]:
            if node not in seen:
                parents[node] = current
                que.append(node)

    result = [None] * len(A)
    for val in tree:
        count = 0
        target = A[val - 1]
        p = val
        while p:
            p = parents[p]
            if p is None: break
            if A[p - 1] > target:
                count += 1
        result[val - 1] = count


    return result

The time complexity of the above algorithm is quadratic in the worst case. It has to run on around 10^6 nodes.

How do I improve the time complexity?

Answer 1

For some graphs complexity of your algorithm is O(n) , eg tree from m nodes there m-1 are direct children of the root node. For some graphs it is O(n^2) , eg tree with one long leg. Then it is 1+2+3+..+(n-1) .

What you can do is instead of iterating through all parents store them in a balanced tree and find position of every next child in logarithmic time instead of your current linear time. With this change your worst time would decrease from quadratic to O(n*log(n)) . See here for a logs summation formula.