Prerequisites :- DFS
Problem :- Given an undirected connected tree with N coloured nodes (colors denoted with integers 1 to N), numbered from 1 to N, and rooted at node 1, determine for each node the number of distinct colors in the subtree of the node.
Solution Approach:
The solution uses Depth-First Search (DFS) to efficiently traverse the tree and maintain a set of distinct colors for each node.
Let’s take a look at the the approach, step by step:
- DFS Traversal:
- Start a DFS traversal from the root node (node 1).
- During the traversal, maintain a set of distinct colors for each node, representing the colors in its subtree.
- Set Operations:
- As the DFS progresses from parent to child nodes, update the set of distinct colors for each node.
- Combine the sets of distinct colors from child nodes with the set of the current node. The main concept is to swap the large and small set and merge the small set to large set always, thus reducing the time complexity by a huge margin.
Time Complexity: O(NlogN + E) = O(NlogN), where N is the number of nodes and E is the no. of edges in the tree. The extra logN factor comes from the set insertion and merging.
Space Complexity: O(N), as the algorithm uses a set and few arrays to store the colors and results.