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 1Input: 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:
#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");
}
}
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
# 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
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 * log2N), where N is the number of nodes in the tree.
Auxiliary space: O(N).
Contact Us