Algorithmic Puzzle: Distinct nodes in a subtree

Question

I am trying to solve this question:

You are given a rooted tree consisting of n nodes. The nodes are numbered 1,2,…,n, and node 1 is the root. Each node has a color.

Your task is to determine for each node the number of distinct colors in the subtree of the node.

The brute force solution is to store a set for each node and them cumulatively merge them in a depth first search. That would run in n^2, not very efficient.

How do I solve this (and the same class of problems) efficiently?

Answer 1

First of all, you'd want to change tree into a list. This technique is often called 'Euler Tour'.

Basically you make an empty list and run DFS. If you visit a node first or last time, push it's color at the end of the list. In this way you'll get list of length 2 * n, where n is equal to number of nodes. It's easy to see that in the list, all colors corresponding to node's children are between its first and last occurrence. Now instead of tree and queries 'how many different colors are there in node's subtree' you have list and queries 'how many different colors are there between index i-th and j-th'. That actually makes things a lot easier.

First idea -- MO's technique O(n sqrt(n)):

I will describe it briefly, I strongly recommend searching up MO's technique, it is well explained in many sources.

Sort all your queries (remainder, they look like this: given pair (i, j) find all distinct numbers in sub-array from index i to index j) by their start. Make sqrt(n) buckets, place query starting from index i to bucket number i / sqrt(n).

For each bucket we will answer the queries separately. Sort all queries in the bucket by their end. Now start processing the first one (the query which end is most to the left) using brute force (iterate over the subarray, store numbers in set/hashset/map/whatever, get size of the set).

Now to process the next one, we shall add some numbers at the end (next query ends farther than the previous one!) and, unfortunately, do something about its start. We'd need to either delete some numbers from the set (if the next query's start > old query start) or add some numbers from the beginning (if the next query's start < old query start). However, we may do it using brute force too, since all queries have start in the same segment of sqrt(n) indices! In total we get O(n sqrt(n)) time complexity.

Second idea -- check this out, O(n log n): Is it possible to query number of distinct integers in a range in O(lg N)?

Answer 2

For each node,

1. Recursively traverse the left and right nodes.
2. Have each call return a HashSet of color.
3. At each node, merge the left child set, the right child set .
4. Update the count for the current node in a HashMap.
5. Add the color of current node and return the set.

---

Sample C# code:

public Dictionary<Node, Integer> distinctColorCount = new ...

public HashSet<Color> GetUniqueColorsTill (TreeNode t) {
    // If null node, return empty set.
    if (t == null) return new HashSet<Color>();

    // If we reached here, we are at a non-null node.
    // First get the set from its left child.
    var lSet = GetUniqueColorsTill(t.Left);

    // Second get the set from its right child.
    var rSet = GetUniqueColorsTill(t.Right);

    // Now, merge the two sets.
    // Can be a little clever here. Merge smaller set to bigger set.
    var returnSet = rSet;
    returnSet.AddAll(lSet);

    // Put the count for this node in the dictionary.
    distinctColorCount[t] = returnSet.Count;    

    // Finally, add the color of current node and return.
    returnSet.Add(t.Color);

    return returnSet;
}

---

You can figure out the complexity exactly as @user58697 commented on your question using the Master Theorem. This is another answer from me written long time ago that explains Master Theorem, if you need a refresher.

c#