HackerRank Unique Colors problem solution YASH PAL, 31 July 2024 In this HackerRank Unique Colors problem, we have given an unrooted tree of n nodes numbered from 1 to n. and we need to print the sum of each node. Problem solution in Java Programming. import java.io.*; import java.util.*; public class Solution { public static int[] sortByPreorder(int n) { int[] stack = new int[n]; int[] ord = new int[n]; boolean[] visited = new boolean[n]; int p = 1; int r = 0; visited[0] = true; while (p > 0) { int cur = stack[p - 1]; ord[r++] = cur; p--; for (int x = ptr[cur]; x > 0; x = nxt[x]) { int e = succ[x]; if (!visited[e]) { stack[p++] = e; visited[e] = true; } } } return ord; } public static int[][] makeRights(int n, int[] par) { int[] ord = sortByPreorder(n); int[] iord = new int[n]; for (int i = 0; i < n; i++) { iord[ord[i]] = i; } int[] right = new int[n]; for (int i = n - 1; i >= 0; i--) { int v = i; for (int x = ptr[ord[i]]; x > 0; x = nxt[x]) { int e = succ[x]; if (e != par[ord[i]]) { v = Math.max(v, right[iord[e]]); } } right[i] = v; } return new int[][] {iord, right}; } public static int[][] parents(int n) { int[] parent = new int[n]; Arrays.fill(parent, -1); int[] queue = new int[n]; for (int p = 0, r = 1; p < r; p++) { int u = queue[p]; for (int i = ptr[u]; i > 0; i = nxt[i]) { int v = succ[i]; if (parent[u] != v) { queue[r++] = v; parent[v] = u; } } } return new int[][] {parent, queue}; } static int[] nxt; static int[] succ; static int[] ptr; static int index = 1; static void addEdge(int u, int v) { nxt[index] = ptr[u]; ptr[u] = index; succ[index++] = v; nxt[index] = ptr[v]; ptr[v] = index; succ[index++] = u; } public static void main(String[] args) throws Exception { BufferedReader br = new BufferedReader(new InputStreamReader(System.in)); BufferedWriter bw = new BufferedWriter(new FileWriter(System.getenv("OUTPUT_PATH"))); StringTokenizer st = new StringTokenizer(br.readLine()); int n = Integer.parseInt(st.nextToken()); int[] a = new int[n]; st = new StringTokenizer(br.readLine()); for (int i = 0; i < n; i++) { int item = Integer.parseInt(st.nextToken()); a[i] = item; } ptr = new int[n]; nxt = new int[2 * n]; succ = new int[2 * n]; for (int i = 0; i < n - 1; i++) { st = new StringTokenizer(br.readLine()); int u = Integer.parseInt(st.nextToken()) - 1; int v = Integer.parseInt(st.nextToken()) - 1; addEdge(u, v); } int[][] pars = parents(n); int[] par = pars[0]; int[] ord = pars[1]; int[][] rs = makeRights(n, par); int[] iord = rs[0]; int[] right = rs[1]; int[] des = new int[n]; for (int i = n - 1; i >= 0; i--) { int cur = ord[i]; des[cur]++; if (i > 0) { des[par[cur]] += des[cur]; } } long[] imos = new long[n + 1]; @SuppressWarnings("unchecked") Map<Integer, List<Integer>>[] dp = new Map[n]; for (int i = n - 1; i >= 0; i--) { int cur = ord[i]; Map<Integer, List<Integer>> map = new HashMap<>(); int color = a[cur]; int ldes = 0; for (int y = ptr[cur]; y > 0; y = nxt[y]) { int e = succ[y]; if (par[cur] != e) { int plus = des[e]; List<Integer> vals = dp[e].get(color); if (vals != null) { for (int val : vals) { plus -= des[val]; } } imos[iord[e]] += plus; imos[right[iord[e]] + 1] -= plus; if (vals != null) { for (int val : vals) { imos[iord[val]] -= plus; imos[right[iord[val]] + 1] += plus; } } Map<Integer, List<Integer>> x = dp[e]; if (ldes < des[e]) { Map<Integer, List<Integer>> temp = x; x = map; map = temp; } for (Map.Entry<Integer, List<Integer>> en : x.entrySet()) { if (!map.containsKey(en.getKey())) { map.put(en.getKey(), new ArrayList<>()); } map.get(en.getKey()).addAll(en.getValue()); } ldes += des[e]; } } map.remove(color); if (i > 0) { List<Integer> sol = new ArrayList<>(); sol.add(cur); map.put(color, sol); dp[cur] = map; } else { for (Map.Entry<Integer, List<Integer>> en : map.entrySet()) { int plus = des[cur]; List<Integer> vals = en.getValue(); for (int val : vals) { plus -= des[val]; } imos[iord[cur]] += plus; imos[right[iord[cur]] + 1] -= plus; for (int val : vals) { imos[iord[val]] -= plus; imos[right[iord[val]] + 1] += plus; } } } } Arrays.sort(a); int uniq = 0; for (int i = 0; i < n; i++) { if (i == 0 || a[i] != a[i - 1]) { uniq++; } imos[i + 1] += imos[i]; } for (int i = 0; i < n; i++) { long res = (long) uniq * n - imos[iord[i]]; bw.write(res + "n"); } bw.newLine(); bw.close(); br.close(); } } Problem solution in C++ programming. #include <cstdio> #include <iostream> #include <ctime> #include <iomanip> #include <cstdlib> #include <algorithm> #include <vector> #include <set> #include <map> #include <string> #include <cstring> using namespace std; vector <int> ve[110000], V[110000]; long long ans[110000], Base, del[410000]; int sum[110000], L[110000], R[110000], st[110000], Time, n, x, y; bool cmp(int x, int y) { return L[x] > L[y]; } void dfs(int k, int f) { L[k] = ++Time; sum[k] = 1; for (int i = 0; i < (int) ve[k].size(); i++) if (ve[k][i] != f) dfs(ve[k][i], k), sum[k] += sum[ve[k][i]]; R[k] = Time; } void modify(int k, int q, int h, int l, int r, int d) { if (l <= q && h <= r) del[k] += d; else { if (r <= (q + h) / 2) modify(k * 2, q, (q + h) / 2, l, r, d); else if ((q + h) / 2 < l) modify(k * 2 + 1, (q + h) / 2 + 1, h, l, r, d); else modify(k * 2, q, (q + h) / 2, l, r, d), modify(k * 2 + 1, (q + h) / 2 + 1, h, l, r, d); } } void dfs(int k, int q, int h, long long now) { now += del[k]; if (q == h) ans[q] = now; else { dfs(k * 2, q, (q + h) / 2, now); dfs(k * 2 + 1, (q + h) / 2 + 1, h, now); } } void doit(vector <int> V) { if (V.size() == 0) return ; Base += n; sort(V.begin(), V.end(), cmp); int len = 0; // printf("xx "); // for (int i = 0; i < (int) V.size(); i++) // printf("%d ", V[i]); // printf("n"); for (int i = 0; i < (int) V.size(); i++) { vector <int> T; while (len && L[V[i]] <= L[st[len]] && L[st[len]] <= R[V[i]]) { T.push_back(st[len]); len--; } st[++len] = V[i]; int r = T.size() - 1; for (int p = ((int) ve[V[i]].size()) - 1; p >= 0; p--) if (sum[ve[V[i]][p]] < sum[V[i]]) { int q = ve[V[i]][p]; int kk = sum[q]; int RR = r; while (RR >= 0 && L[q] <= L[T[RR]] && L[T[RR]] <= R[q]) RR -= 1; for (int j = RR + 1; j <= r; j++) kk -= sum[T[j]]; // kk -= 1; // printf("%d %dn", V[i], kk); // if (L[V[i]] != R[V[i]]) modify(1, 1, n, L[q], R[q], kk); for (int j = RR + 1; j <= r; j++) modify(1, 1, n, L[T[j]], R[T[j]], -kk); r = RR; } } // for (int i = 1; i <= len; i++) // printf("%d ", st[i]); // printf("n"); int kk = n; for (int j = 1; j <= len; j++) kk -= sum[st[j]]; modify(1, 1, n, 1, n, kk); for (int j = 1; j <= len; j++) modify(1, 1, n, L[st[j]], R[st[j]], -kk); } int main() { scanf("%d", &n); for (int i = 1; i <= n; i++) scanf("%d", &x), V[x].push_back(i); for (int i = 1; i < n; i++) { scanf("%d%d", &x, &y); ve[x].push_back(y); ve[y].push_back(x); } dfs(1, 0); for (int i = 1; i <= 100000; i++) { doit(V[i]); } dfs(1, 1, n, 0); for (int i = 1 ; i <= n; i++) printf("%lldn", Base - ans[L[i]]); } Problem solution in C programming. #include <stdio.h> #include <string.h> void descending_order(unsigned *self, unsigned *weights, unsigned cnt) { unsigned at, member, node = cnt >> 1; for (self--; node; self[at >> 1] = member) for (member = self[node], at = node-- << 1; at <= cnt; at <<= 1) { at |= (at < cnt) && (weights[self[at | 1]] < weights[self[at]]); if (weights[member] <= weights[self[at]]) break ; self[at >> 1] = self[at]; } for (; cnt > 1; self[at >> 1] = member) { member = self[cnt]; self[cnt--] = self[1]; for (at = 2; at <= cnt; at <<= 1) { at |= (at < cnt) && (weights[self[at | 1]] < weights[self[at]]); if (weights[member] < weights[self[at]]) break ; self[at >> 1] = self[at]; } } } #define have(self, id) ((self)[(id) >> 5] & (1U << ((id) & 31U))) #define inv(self, id) ((self)[(id) >> 5] ^= (1U << ((id) & 31U))) int main() { unsigned at, tail, others, vertex_cnt; scanf("%u", &vertex_cnt); unsigned colors[vertex_cnt], ancestors[vertex_cnt], centroids[vertex_cnt], history[vertex_cnt], descendants[vertex_cnt << 1], indices[vertex_cnt + 1], weights[vertex_cnt + 1], distinct = 0; for (at = 0; at < vertex_cnt; at++) if (distinct < (scanf("%u", &colors[at]), (colors[at] - 1))) distinct = colors[at] - 1; { unsigned seen[((distinct + 1) >> 5) + 1]; for (memset(seen, 0, sizeof(seen)); at--; colors[at] = history[colors[at]]) if (have(seen, colors[at]) == 0) { inv(seen, colors[at]); history[colors[at]] = distinct++; } } for (at = vertex_cnt >> 1; at; ((unsigned long *)ancestors)[--at] = 0x100000001UL * vertex_cnt); for (ancestors[vertex_cnt - 1] = vertex_cnt; ++at < vertex_cnt; ancestors[others - 1] = tail - 1) if (ancestors[scanf("%u %u", &others, &tail), others - 1] != vertex_cnt) others ^= tail, tail ^= others, others ^= tail; for (memset(indices, 0, sizeof(indices)); at--; indices[ancestors[at]]++); for (; ++at < vertex_cnt; indices[at + 1] += indices[at]); for (; at--; descendants[--indices[ancestors[at]]] = at); for (others = 0, history[at = vertex_cnt - 1] = descendants[vertex_cnt - 1]; others < at; others++) { history[others] = history[at]; at -= (indices[history[at] + 1] - indices[history[at]]) - 1; memcpy( &history[at], &descendants[indices[history[others]]], sizeof(history[0]) * (indices[history[others] + 1] - indices[history[others]]) ); } for (at = vertex_cnt >> 1; at--; ((unsigned long *)weights)[at] = 0x100000001UL); for (weights[(at = vertex_cnt) - 1] = 1; --at; weights[ancestors[history[at]]] += weights[history[at]]); for (others = indices[at = history[0]]; others < indices[at + 1]; others++) if ((weights[descendants[others]] << 1) > vertex_cnt) { weights[at] -= weights[descendants[others]]; weights[descendants[others]] += weights[at]; ancestors[at] = descendants[others]; others = indices[at = descendants[others]] - 1; } memset(indices, 0, sizeof(indices)); for (ancestors[at] = vertex_cnt, at = vertex_cnt; at--; indices[ancestors[at]]++); for (; ++at < vertex_cnt; indices[at + 1] += indices[at]); for (; at--; descendants[--indices[ancestors[at]]] = at); for (; ++at < vertex_cnt; descending_order(&descendants[indices[at]], weights, indices[at + 1] - indices[at])); unsigned root; { unsigned prev, mass[at--]; history[at] = descendants[at]; centroids[history[at]] = (mass[at] = vertex_cnt); weights[vertex_cnt] = 0; for (tail = 0; tail < at; weights[root] = 0) { for (root = (prev = history[at]); (weights[root] << 1) < mass[at]; root = ancestors[prev = root]); for (others = indices[root]; others < indices[root + 1]; others++) if ((weights[descendants[others]] << 1) > mass[at] && descendants[others] != prev) others = indices[root = descendants[others]] - 1; centroids[history[tail++] = root] = centroids[history[at++]]; for (others = root; weights[ancestors[others]]; weights[others = ancestors[others]] -= weights[root]); if (others != root) { mass[--at] = weights[others]; centroids[history[at] = ancestors[root]] = root; } for (others = indices[root]; others < indices[root + 1]; others++) if (weights[descendants[others]]) { mass[--at] = weights[descendants[others]]; centroids[history[at] = descendants[others]] = root; } } } unsigned centroids_indices[vertex_cnt + 1], centroids_descendants[vertex_cnt]; for (memset(centroids_indices, 0, sizeof(centroids_indices)), at = vertex_cnt; at--; centroids_indices[centroids[at]]++); for (; ++at < vertex_cnt; centroids_indices[at + 1] += centroids_indices[at]); for (; at--; centroids_descendants[--centroids_indices[centroids[at]]] = at); // submitted by Samy Vilar <samy_vilar> 09/20/2018 for (at = vertex_cnt >> 1; at--; ((unsigned long *)indices)[at] = 0x100000001UL); for (*(unsigned long *)&indices[(at = vertex_cnt) - 1] = 1UL; at--; indices[ancestors[at]]++); for (; ++at < vertex_cnt; indices[at + 1] += indices[at]); for (; at--; descendants[--indices[ancestors[at]]] = at); for (; ++at < vertex_cnt; descendants[--indices[at]] = ancestors[at]); unsigned long totals[vertex_cnt]; unsigned successors[vertex_cnt], local[vertex_cnt + 1], freqs[distinct], unique[distinct], trace[distinct], global[distinct], base_colors[distinct], seen[((vertex_cnt + 1) >> 5) + 1]; memset(totals, 0, sizeof(totals)); memset(global, 0, sizeof(global)); memset(freqs, 0, sizeof(freqs)); memset(seen, 0, sizeof(seen)); inv(seen, vertex_cnt); // Samy Vilar <samy_vilar> 10/20/2018 for (at = distinct >> 1; at--; ((unsigned long *)trace)[at] = 0x100000001UL * vertex_cnt); trace[distinct - 1] = vertex_cnt; local[vertex_cnt] = 0; distinct = 0; for (at = 1; at--; freqs[colors[root]] = 0) { root = history[tail = at]; inv(seen, root); for (others = indices[root]; others < indices[root + 1]; others++) if (have(seen, descendants[others]) == 0) history[at++] = descendants[others]; unsigned branch_cnt = at - tail, base_cnt = 0, curr; unsigned long aids[branch_cnt]; for (memset(aids, 0, sizeof(aids)); distinct; global[unique[--distinct]] = 0); unsigned long total = 0; for (freqs[colors[root]] = (weights[root] = 1); branch_cnt--; weights[root] += weights[history[++at]]) { while (at-- != (tail + branch_cnt)) if (inv(seen, history[at]) & (1U << (history[at] & 31U))) { weights[curr = history[at++]] = 1; freqs[colors[curr]]++; for (others = indices[curr]; others < indices[curr + 1]; others++) if (have(seen, descendants[others]) == 0) history[at++] = descendants[others]; } else { for (others = indices[history[at]]; others < indices[history[at] + 1]; others++) if (have(seen, descendants[others]) == 0) weights[history[at]] += weights[descendants[others]]; if (--freqs[colors[history[at]]] == 0) { if (trace[colors[history[at]]] == vertex_cnt) base_colors[base_cnt++] = colors[history[at]]; successors[history[at]] = trace[colors[history[at]]]; trace[colors[history[at]]] = history[at]; local[history[at]] = local[successors[history[at]]] + weights[history[at]]; } } for (; base_cnt; trace[base_colors[base_cnt]] = vertex_cnt) { for (others = successors[trace[base_colors[--base_cnt]]]; others != vertex_cnt; others = successors[others]) local[others] = local[trace[base_colors[base_cnt]]]; if (global[base_colors[base_cnt]] == 0) unique[distinct++] = base_colors[base_cnt]; global[base_colors[base_cnt]] += local[trace[base_colors[base_cnt]]]; aids[branch_cnt] += local[trace[base_colors[base_cnt]]]; } total += aids[branch_cnt]; } totals[root] += total + weights[root]; for (others = indices[root]; others < indices[root + 1]; others++) if (have(seen, descendants[others]) == 0) history[at++] = descendants[others]; unsigned dist = 1; for (branch_cnt = at - tail; branch_cnt--; total += aids[branch_cnt], at++) for (total -= aids[branch_cnt]; at-- != (tail + branch_cnt); ) if (inv(seen, history[at]) & (1U << (history[at] & 31U))) { if (freqs[colors[history[at]]]++ == 0) total -= global[colors[history[at]]] - local[history[at]], dist++; totals[history[at]] += total + (dist * (weights[root] - weights[history[tail + branch_cnt]])); for (others = indices[curr = history[at++]]; others < indices[curr + 1]; others++) if (have(seen, descendants[others]) == 0) history[at++] = descendants[others]; } else if (--freqs[colors[history[at]]] == 0) total += global[colors[history[at]]] - local[history[at]], dist--; memcpy( &history[at], ¢roids_descendants[centroids_indices[root]], sizeof(history[0]) * (centroids_indices[root + 1] - centroids_indices[root]) ); at += centroids_indices[root + 1] - centroids_indices[root]; } for (; ++at < vertex_cnt; printf("%lun", totals[at])); return 0; } coding problems data structure