算法难题 : Distinct nodes in a subtree

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.

强力解决方案是为每个节点存储一个集合，并在深度优先搜索中累积合并它们。这将在 n^2 中运行，效率不高。

如何有效地解决这个问题(以及同类问题)？

Answer 1

对于每个节点，

1. 递归遍历左右节点。
2. 让每个调用返回一个颜色的哈希集。
3. 在每个节点，合并左子集、右子集。
4. 更新 HashMap 中当前节点的计数。
5. 添加当前节点的颜色并返回集合。

---

示例 C# 代码:

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;
}

---

您可以使用主定理准确地计算出复杂性，就像 @user58697 对您的问题的评论一样。 This这是我很久以前写的另一个答案，它解释了主定理，如果您需要复习的话。

c#