Problem statement

https://binarysearch.com/problems/Popularity/

Solution

For each node we need to calculate number of children we have. For edge x -> y, where y lies below x, we need to calculate dp[y] * (n - dp[y]).

Complexity

It is O(n) for time and space.

Code

class Solution:
    def solve(self, edges):
        G = defaultdict(list)
        n = len(edges) + 1
        d = {}
        for i, (x, y) in enumerate(edges):
            G[x] += [y]
            G[y] += [x]
            d[x, y] = i

        dp = {}
        def dfs(node, par):
            if not G[node]: return 0
            ans = 1
            for child in G[node]:
                if child == par: continue
                ans += dfs(child, node)
            dp[node] = ans
            return ans

        dfs(0, -1)
        ans = []
        for x, y in edges:
            t = min(dp[x], dp[y])
            ans += [t * (n - t)]

        return ans