Official Analysis (C++)

Explanation

Since this problem involves the answering of a series of subtree queries, we should probably use a Euler tour of some sort to solve this problem. For the rest of this problem, $\texttt{start}[x]$ will be the first index at which the snowball appears in the tour, and $\texttt{end}[x]$ will be the second (and last) index at which the snowball appears in the tour.

There's a couple of problems we have to address for each query:

1. Detecting when a snowball has already been colored a certain color. If we detect this, the query is meaningless and won't change anything, since the parent's coloring already added the colors to all the subtree snowballs.
2. Detecting when a snowball has children that have already been colored a certain color. In this case, we should probably remove those children from consideration, since the current query will cover all of them.
3. Actually calculating the colorfulness of a subtree.

Issues 1 & 2

The first two problems are solvable with our Euler tour arrays. For each color, we create a sorted map with the keys being $\texttt{start}[x]$ and the values being $x$, where $x$ represents a snowball. When we add a new snowball, we can use the map's lower key and upper key functions to efficiently remove any children or detect a parent.

Issue 3

After addressing those two problems, we have a map for each color, where no two relevant snowballs of a color are either parents or children of each other. Keep in mind that this only applies to each individual map: a snowball in color $c_1$'s map can still have a child in color $c_2$'s map.

For each part, we can use a separate point update range query data structure like a binary indexed tree or segment tree.

Parent Colorfulness

This is the amount resulting from the parents' and current snowball's paint.

When adding a snowball to our color map, we increment its index at $\texttt{start}[x]$ by $1$ and decrement its index at $\texttt{end}[x]$ by 1. When removing it, we do the reverse. Now, if we perform a range query of the sum of all numbers from $0$ to $\texttt{start}[x]$, we can get the number of colors that result from parents. However, since these colors are applied to every snowball in the subtree, we need to multiply this by the subtree size of the current snowball as well.

Child Colorfulness

This is the amount resulting from children's paint.

For this part, when adding a snowball, we add the current snowball's subtree size to $\texttt{start}[x]$ (and do the reverse when removing a snowball). Then, by performing a range query of the sum of all values from $\texttt{start}[x] + 1$ to $\texttt{end}[x]$, we get the subtree sizes of all the childrens' unique colors.

By adding these two values together, we get our answer for the query.

Implementation

Time Complexity: $\mathcal{O}(N+Q\log N)$

Copy
#include <fstream>
#include <iostream>
#include <map>
#include <vector>

using std::cout;
using std::endl;
using std::vector;

 Code Snippet: Binary Indexed Tree (Click to expand)

Copy
import java.io.*;
import java.util.*;

public class SnowCow {
	public static void main(String[] args) throws IOException {
		BufferedReader read = new BufferedReader(new FileReader("snowcow.in"));
		StringTokenizer initial = new StringTokenizer(read.readLine());
		int snowballNum = Integer.parseInt(initial.nextToken());
		int queryNum = Integer.parseInt(initial.nextToken());