CSES Solutions – Distinct Colors

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.

Example:

Input: n = 5, color_input = {2, 3, 2, 2, 1}, edges = {{1, 2}, {1, 3}, {3, 4}, {3, 5}}
Output: 3 1 2 1 1

Input: n = 5, color_input = {2, 3, 2, 1, 1}, edges = {{2, 3}, {1, 3}, {3, 4}, {3, 5}}
Output: 3 1 3 1 1

Approach:

The solution involves a clever optimization to efficiently merge sets of colors during the tree traversal.

Here’s a breakdown of the key points:

• Merging Sets: The naive approach to merging two sets of colors involves iterating through one set and inserting each element into the other set. This takes O(m * log(n + m)) time, where m and n are the sizes of the sets.
• Optimization: The solution optimizes this process by ensuring that the smaller set is always merged into the larger set. This is achieved using the swap function, which exchanges the two sets in O(1) time.
• Time Complexity Analysis: By merging the smaller set into the larger one, the size of the resulting set at least doubles. This means that an element that has been moved Y times will be in a set of size at least 2^Y. Since the maximum size of a set is N (the root), each element will be moved at most O(log N) times. Therefore, the overall time complexity of the solution is O(N * log^2 N).

Steps to solve this problem:

Explanation of the Function:

• Maintain a function process_colors to iterate over children of current node curr, and for each child node
  • Call process_colors recursively to process the subtree of n.
  • If the color set of n is larger than the color set of curr, swap them to ensure the larger set is always in colors[curr].
  • Merge the color sets of n and curr by inserting all elements from colors[n] into colors[curr].
• Update distinct colors:
  • Set the value in distinct_num[curr] to the size of the colors[curr] set, representing the number of distinct colors in the subtree of curr.

Below is the implementation of above approach:

C++

#include <bits/stdc++.h>
using namespace std;

// Maximum number of nodes
const int MAX_N = 2e5;

// Adjacency list to store the tree
vector<int> adj[MAX_N + 1];

// Set to store distinct colors at each node
set<int> colors[MAX_N + 1];

// Array to store the number of distinct colors at each node
int distinct_num[MAX_N + 1];

// Function to process colors in the tree
void process_colors(int curr, int parent)
{
    // Iterate over all children of the current node
    for (int n : adj[curr]) {
        // Avoid processing the parent node
        if (n != parent) {
            // Recursively process child nodes
            process_colors(n, curr);
            // Ensure the larger set is always in
            // colors[curr]
            if (colors[curr].size() < colors[n].size()) {
                swap(colors[curr], colors[n]);
            }
            // Merge the color sets
            for (int item : colors[n]) {
                colors[curr].insert(item);
            }
        }
    }
    // Update the number of distinct colors at the current
    // node
    distinct_num[curr] = colors[curr].size();
}

int main()
{
    // Static input
    // Number of nodes
    int n = 5;
    // Node colors
    vector<int> color_input = { 2, 3, 2, 2, 1 };
    // Edges
    vector<pair<int, int> > edges
        = { { 1, 2 }, { 1, 3 }, { 3, 4 }, { 3, 5 } };

    // Assign colors to nodes
    for (int i = 1; i <= n; i++) {
        colors[i].insert(color_input[i - 1]);
    }

    // Construct the tree using the edges
    for (auto& edge : edges) {
        int a = edge.first, b = edge.second;
        adj[a].push_back(b);
        adj[b].push_back(a);
    }

    // Process colors starting from the root node (1)
    process_colors(1, 0);

    // Output the number of distinct colors for each node
    for (int i = 1; i <= n; i++) {
        cout << distinct_num[i] << (i < n ? " " : "\n");
    }
}

Java

import java.util.*;

public class ColorProcessingTree {
    // Maximum number of nodes
    static final int MAX_N = 200000;

    // Adjacency list to store the tree
    static List<Integer>[] adj = new ArrayList[MAX_N + 1];

    // Set to store distinct colors at each node
    static Set<Integer>[] colors = new HashSet[MAX_N + 1];

    // Array to store the number of distinct colors at each
    // node
    static int[] distinctNum = new int[MAX_N + 1];

    // Function to process colors in the tree
    static void processColors(int curr, int parent)
    {
        for (int n : adj[curr]) {
            if (n != parent) {
                processColors(n, curr);
                if (colors[curr].size()
                    < colors[n].size()) {
                    Set<Integer> temp = colors[curr];
                    colors[curr] = colors[n];
                    colors[n] = temp;
                }
                colors[curr].addAll(colors[n]);
            }
        }
        distinctNum[curr] = colors[curr].size();
    }

    public static void main(String[] args)
    {
        // Static input
        int n = 5;
        int[] colorInput = { 2, 3, 2, 2, 1 };
        int[][] edges
            = { { 1, 2 }, { 1, 3 }, { 3, 4 }, { 3, 5 } };

        // Initialize adjacency list and color sets
        for (int i = 1; i <= n; i++) {
            adj[i] = new ArrayList<>();
            colors[i] = new HashSet<>();
            colors[i].add(colorInput[i - 1]);
        }

        // Construct the tree using the edges
        for (int[] edge : edges) {
            int a = edge[0], b = edge[1];
            adj[a].add(b);
            adj[b].add(a);
        }

        // Process colors starting from the root node (1)
        processColors(1, 0);

        // Output the number of distinct colors for each
        // node
        for (int i = 1; i <= n; i++) {
            System.out.print(distinctNum[i]);
            if (i < n) {
                System.out.print(" ");
            }
            else {
                System.out.println();
            }
        }
    }
}

// This code is contributed by shivamgupta0987654321

Python3

# Function to process colors in the tree
def process_colors(curr, parent):
    for n in adj[curr]:
        if n != parent:
            process_colors(n, curr)
            if len(colors[curr]) < len(colors[n]):
                colors[curr], colors[n] = colors[n], colors[curr]
            colors[curr] |= colors[n]
    distinct_num[curr] = len(colors[curr])


# Static input
n = 5
color_input = [2, 3, 2, 2, 1]
edges = [(1, 2), (1, 3), (3, 4), (3, 5)]

# Assign colors to nodes
colors = [set() for _ in range(n + 1)]
for i in range(1, n + 1):
    colors[i].add(color_input[i - 1])

# Construct the tree using the edges
adj = [[] for _ in range(n + 1)]
for a, b in edges:
    adj[a].append(b)
    adj[b].append(a)

# Initialize distinct_num array
distinct_num = [0] * (n + 1)

# Process colors starting from the root node (1)
process_colors(1, 0)

# Output the number of distinct colors for each node
print(*distinct_num[1:])

# This code is contributed by Ayush Mishra

JavaScript

class ColorProcessingTree {
    // Maximum number of nodes
    static MAX_N = 200000;

    constructor() {
        // Adjacency list to store the tree
        this.adj = Array.from({ length: ColorProcessingTree.MAX_N + 1 }, () => []);

        // Set to store distinct colors at each node
        this.colors = Array.from({ length: ColorProcessingTree.MAX_N + 1 }, () => new Set());

        // Array to store the number of distinct colors at each node
        this.distinctNum = Array.from({ length: ColorProcessingTree.MAX_N + 1 }, () => 0);
    }

    // Function to process colors in the tree
    processColors(curr, parent) {
        for (const n of this.adj[curr]) {
            if (n !== parent) {
                this.processColors(n, curr);
                if (this.colors[curr].size < this.colors[n].size) {
                    const temp = this.colors[curr];
                    this.colors[curr] = this.colors[n];
                    this.colors[n] = temp;
                }
                this.colors[curr] = new Set([...this.colors[curr], ...this.colors[n]]);
            }
        }
        this.distinctNum[curr] = this.colors[curr].size;
    }

    main() {
        // Static input
        const n = 5;
        const colorInput = [2, 3, 2, 2, 1];
        const edges = [[1, 2], [1, 3], [3, 4], [3, 5]];

        // Initialize adjacency list and color sets
        for (let i = 1; i <= n; i++) {
            this.colors[i].add(colorInput[i - 1]);
        }

        // Construct the tree using the edges
        for (const edge of edges) {
            const [a, b] = edge;
            this.adj[a].push(b);
            this.adj[b].push(a);
        }

        // Process colors starting from the root node (1)
        this.processColors(1, 0);

        // Output the number of distinct colors for each node
        for (let i = 1; i <= n; i++) {
            process.stdout.write(`${this.distinctNum[i]}`);
            if (i < n) {
                process.stdout.write(" ");
            } else {
                process.stdout.write("\n");
            }
        }
    }
}

// Instantiate and run the main method
const colorProcessingTree = new ColorProcessingTree();
colorProcessingTree.main();

Output

3 1 2 1 1

Time complexity: O(N * log²N), where N is the number of nodes in the tree.
Auxiliary space: O(N).