Heavy-Light Chain Count

Problem

Count the heavy chains built by a heavy-light decomposition of a tree rooted at 0.

Compute subtree sizes, pick for each node the child with the largest subtree as its heavy child (ties by smallest index), then count one chain per node that is either the root or not the heavy child of its parent. edges is a list of [u, v] pairs. Return the number of chains.

Hint

A chain starts exactly at a light edge, so count the light edges plus the root's chain.

Test cases

case 1

in: [4,[[0,1],[1,2],[2,3]],0]

expected: 1

case 2

in: [4,[[0,1],[0,2],[0,3]],0]

expected: 3

case 3

in: [6,[[0,1],[0,2],[1,3],[1,4],[2,5]],0]

expected: 3

case 4

in: [5,[[0,1],[1,2],[2,3],[3,4]],0]

expected: 1

case 5

in: [1,[],0]

expected: 1

Starter codePython

def heavy_light_chain_count(n, edges, root):
    # Your code here
    pass

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