Official Analysis

Implementation

As explained in the official analysis, when we use small to large merging on data structures such as sets and maps, we optimize the amount of operations from $\mathcal{O}(N^2 \log N)$ to $\mathcal{O}(N \log^2 N)$. We can take advantage of this to store all kinds of information related to colors, their sum, and their occurrences.

Time Complexity: $\mathcal{O}(N \log^2 N)$

Copy
#include <bits/stdc++.h>
using namespace std;
using ll = long long;

const int MAX_N = 1e5;

array<ll, MAX_N> color;
array<ll, MAX_N> ans;
vector<int> g[MAX_N];