Heavy-Light Decomposition Head Selection

Problem

Compute heavy-light decomposition heads for a rooted tree.

parents[i] is the parent of node i (-1 for the root). Subtree sizes determine each node's heavy child: the largest child subtree, with ties going to the smallest index. A light child starts a new chain, while a heavy child inherits its parent's head.

Return [node, heavy_child, head] triples for every node in index order.

Hint

Every light edge starts a new chain and heavy edges continue the parent's chain.

Test cases

case 1

in: [[-1,0,0,1,1,2]]

expected: [[0,1,0],[1,3,0],[2,5,2],[3,-1,0],[4,-1,4],[5,-1,2]]

case 2

in: [[-1]]

expected: [[0,-1,0]]

case 3

in: [[-1,0]]

expected: [[0,1,0],[1,-1,0]]

case 4

in: [[-1,0,1]]

expected: [[0,1,0],[1,2,0],[2,-1,0]]

Starter codePython

def hld_heads(parents):
    # Your code here
    pass

4 test cases · run your code against them in the editor