Bessie's Snow Cow

Explicación

Como este problema involucra responder una serie de consultas de subárbol, probablemente deberíamos usar un tour de Euler de algún tipo para resolverlo. Para el resto de este problema, $\texttt{start}[x]$ será el primer índice en el que la bola de nieve aparece en el tour, y $\texttt{end}[x]$ será el segundo (y último) índice en el que la bola de nieve aparece en el tour.

Hay un par de problemas que tenemos que abordar para cada consulta:

Detectar cuándo una bola de nieve ya fue coloreada de un cierto color. Si detectamos esto, la consulta no tiene sentido y no cambiará nada, ya que el coloreado del padre ya agregó los colores a todas las bolas de nieve del subárbol.
Detectar cuándo una bola de nieve tiene hijos que ya fueron coloreados de un cierto color. En este caso, probablemente deberíamos quitar esos hijos de consideración, ya que la consulta actual los cubrirá a todos.
Calcular realmente la colorful-ness de un subárbol.

Problemas 1 y 2

Los primeros dos problemas se pueden resolver con nuestros arreglos del tour de Euler. Para cada color, creamos un mapa ordenado cuyas claves son $\texttt{start}[x]$ y cuyos valores son $x$ , donde $x$ representa una bola de nieve. Cuando agregamos una nueva bola de nieve, podemos usar las funciones de clave inferior y clave superior del mapa para quitar de forma eficiente cualquier hijo o detectar un padre.

Problema 3

Después de abordar esos dos problemas, tenemos un mapa para cada color, donde no hay dos bolas de nieve relevantes de un color que sean padre o hijo una de la otra. Hay que tener en cuenta que esto solo aplica a cada mapa individual: una bola de nieve en el mapa del color $c_1$ aún puede tener un hijo en el mapa del color $c_2$ .

Para cada parte, podemos usar una estructura de datos separada de actualización puntual y consulta de rango como un Árbol de Fenwick o un Árbol de Segmentos.

Colorful-ness de los padres

Esta es la cantidad que resulta de la pintura de los padres y de la bola de nieve actual.

Al agregar una bola de nieve a nuestro mapa de color, incrementamos su índice en $\texttt{start}[x]$ en $1$ y decrementamos su índice en $\texttt{end}[x]$ en 1. Al quitarla, hacemos lo inverso. Ahora, si hacemos una consulta de rango de la suma de todos los números de $0$ a $\texttt{start}[x]$ , podemos obtener el número de colores que resultan de los padres. Sin embargo, como estos colores se aplican a cada bola de nieve del subárbol, también necesitamos multiplicar esto por el tamaño del subárbol de la bola de nieve actual.

Colorful-ness de los hijos

Esta es la cantidad que resulta de la pintura de los hijos.

Para esta parte, al agregar una bola de nieve, sumamos el tamaño del subárbol de la bola de nieve actual a $\texttt{start}[x]$ (y hacemos lo inverso al quitar una bola de nieve). Luego, al hacer una consulta de rango de la suma de todos los valores de $\texttt{start}[x] + 1$ a $\texttt{end}[x]$ , obtenemos los tamaños de subárbol de todos los colores únicos de los hijos.

Al sumar estos dos valores, obtenemos nuestra respuesta para la consulta.

Implementación

Complejidad temporal: $\mathcal{O}(N+Q\log N)$

C++Java


#include <fstream>
#include <iostream>
#include <map>
#include <vector>
 
using std::cout;
using std::endl;
using std::vector;
 
// BeginCodeSnip{Binary Indexed Tree}
class BITree {
  private:
	vector<long long> bit;
	int size;
 
  public:
	BITree(int size) : size(size), bit(size) {}
 
	void increment(int ind, long long val) {
		ind++;  // have the driver code not worry about 1-indexing
		for (; ind <= size; ind += ind & -ind) { bit[ind] += val; }
	}
 
	long long query(int ind) {
		ind++;
		long long sum = 0;
		for (; ind > 0; ind -= ind & -ind) { sum += bit[ind]; }
		return sum;
	}
};
// EndCodeSnip
 
int main() {
	std::ifstream read("snowcow.in");
	int snowball_num;
	int query_num;
	read >> snowball_num >> query_num;
 
	vector<vector<int>> neighbors(snowball_num);
	for (int b = 0; b < snowball_num - 1; b++) {
		int sb1, sb2;
		read >> sb1 >> sb2;
		sb1--;
		sb2--;
		neighbors[sb1].push_back(sb2);
		neighbors[sb2].push_back(sb1);
	}
 
	// perform our euler tour- start & end are as described previously
	vector<bool> processed(snowball_num);
	vector<int> start(snowball_num);
	vector<int> end(snowball_num);
	int timer = 0;
	vector<int> frontier{0};
	while (!frontier.empty()) {
		int curr = frontier.back();
		frontier.pop_back();
		if (processed[curr]) {
			end[curr] = timer;
			timer++;
			continue;
		}
		start[curr] = timer;
		frontier.push_back(curr);  // set a marker to record the outtime
		processed[curr] = true;
		for (int n : neighbors[curr]) {
			if (!processed[n]) { frontier.push_back(n); }
		}
		timer++;
	}
 
	// calculate subtree sizes of all snowballs
	vector<int> sub_size(snowball_num);
	for (int s = 0; s < snowball_num; s++) {
		sub_size[s] = (end[s] - start[s] + 1) / 2;
	}
 
	// the map for each color for tracking of parents & children
	std::map<int, std::map<int, int>> colors;
	// BIT for calculating the colorfulness due to parents & current snowball
	BITree above_unique(snowball_num * 2);
	// and for calculating the colorfulness due to colors of children
	BITree below_unique(snowball_num * 2);
	std::ofstream written("snowcow.out");
	for (int q = 0; q < query_num; q++) {
		int type;
		int sb;
		read >> type >> sb;
		sb--;
		if (type == 1) {
			int color;
			read >> color;
 
			std::map<int, int> &painted = colors[color];
			auto l_closest = painted.upper_bound(start[sb]);
			/*
			 * check if there's a parent of the snowball that's
			 * already been painted with the color
			 */
			if (l_closest != painted.begin() && end[sb] <= end[(--l_closest)->second]) {
				continue;
			}
			// remove all children of the current snowball that are in the map
			while (true) {
				auto r_next = painted.upper_bound(start[sb]);
				if (r_next == painted.end() || end[sb] <= end[r_next->second]) {
					break;
				}
				int r_next_sb = r_next->second;
				// undo the increments from when the child was added
				above_unique.increment(start[r_next_sb], -1);
				above_unique.increment(end[r_next_sb], 1);
				below_unique.increment(start[r_next_sb], -sub_size[r_next_sb]);
				painted.erase(r_next->first);
			}
 
			painted[start[sb]] = sb;
			above_unique.increment(start[sb], 1);
			above_unique.increment(end[sb], -1);
			below_unique.increment(start[sb], sub_size[sb]);
		} else if (type == 2) {
			/*
			 * the colors from parents influence every node in the subtree,
			 * so multiply by the current subtree size
			 */
			long long above_colors = sub_size[sb] * above_unique.query(start[sb]);
			long long below_colors =
			    (below_unique.query(end[sb]) - below_unique.query(start[sb]));
			written << above_colors + below_colors << '\n';
		}
	}
}


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());
 
		List<Integer>[] neighbors = new ArrayList[snowballNum];
		// sb short for SnowBall
		for (int sb = 0; sb < snowballNum; sb++) { neighbors[sb] = new ArrayList<>(); }
		for (int b = 0; b < snowballNum - 1; b++) {
			StringTokenizer branch = new StringTokenizer(read.readLine());
			int sb1 = Integer.parseInt(branch.nextToken()) - 1;
			int sb2 = Integer.parseInt(branch.nextToken()) - 1;
			neighbors[sb1].add(sb2);
			neighbors[sb2].add(sb1);
		}
 
		// perform our euler tour- start & end are as described previously
		boolean[] processed = new boolean[snowballNum];
		int[] start = new int[snowballNum];
		int[] end = new int[snowballNum];
		int timer = 0;
		ArrayDeque<Integer> frontier = new ArrayDeque<>();
		frontier.add(0);
		while (!frontier.isEmpty()) {
			int curr = frontier.removeLast();
			if (processed[curr]) {
				end[curr] = timer;
				timer++;
				continue;
			}
			start[curr] = timer;
			frontier.add(curr);  // set a marker to record the outtime
			processed[curr] = true;
			for (int n : neighbors[curr]) {
				if (!processed[n]) { frontier.add(n); }
			}
			timer++;
		}
 
		// calculate subtree sizes of all snowballs
		int[] subSize = new int[snowballNum];
		for (int s = 0; s < snowballNum; s++) {
			subSize[s] = (end[s] - start[s] + 1) / 2;
		}
 
		// the map for each color for tracking of parents & children
		Map<Integer, TreeMap<Integer, Integer>> colors = new HashMap<>();
		// BIT for calculating the colorfulness due to parents & current
		// snowball
		BITree aboveUnique = new BITree(snowballNum * 2);
		// and for calculating the colorfulness due to colors of children
		BITree belowUnique = new BITree(snowballNum * 2);
		StringBuilder ans = new StringBuilder();
		for (int q = 0; q < queryNum; q++) {
			StringTokenizer query = new StringTokenizer(read.readLine());
			int type = Integer.parseInt(query.nextToken());
			int sb = Integer.parseInt(query.nextToken()) - 1;
			if (type == 1) {
				int color = Integer.parseInt(query.nextToken());
				if (!colors.containsKey(color)) { colors.put(color, new TreeMap<>()); }
 
				TreeMap<Integer, Integer> painted = colors.get(color);
				Integer lClosest = painted.floorKey(start[sb]);
				/*
				 * check if there's a parent of the snowball that's
				 * already been painted with the color
				 */
				if (lClosest != null && end[sb] <= end[painted.get(lClosest)]) {
					continue;
				}
				// remove all children of the current snowball that are in the
				// map
				while (true) {
					Integer rNext = painted.higherKey(start[sb]);
					if (rNext == null || end[sb] <= end[painted.get(rNext)]) { break; }
					int rNextSB = painted.get(rNext);
					// undo the increments from when the child was added
					aboveUnique.increment(start[rNextSB], -1);
					aboveUnique.increment(end[rNextSB], 1);
					belowUnique.increment(start[rNextSB], -subSize[rNextSB]);
					painted.remove(rNext);
				}
 
				painted.put(start[sb], sb);
				aboveUnique.increment(start[sb], 1);
				aboveUnique.increment(end[sb], -1);
				belowUnique.increment(start[sb], subSize[sb]);
			} else if (type == 2) {
				/*
				 * the colors from parents influence every node in the subtree,
				 * so multiply by the current subtree size
				 */
				long aboveColors = subSize[sb] * aboveUnique.query(start[sb]);
				long belowColors =
				    (belowUnique.query(end[sb]) - belowUnique.query(start[sb]));
				ans.append(aboveColors + belowColors).append('\n');
			}
		}
 
		PrintWriter written = new PrintWriter("snowcow.out");
		written.print(ans);
		written.close();
	}
}
 
// BeginCodeSnip{Binary Indexed Tree}
class BITree {
	private final long[] bit;
	private final int size;
 
	public BITree(int size) {
		bit = new long[size + 1];
		this.size = size;
	}
 
	public void increment(int ind, long val) {
		ind++;  // have the driver code not worry about 1-indexing
		for (; ind <= size; ind += ind & -ind) { bit[ind] += val; }
	}
 
	public long query(int ind) {  // the bound's inclusive
		ind++;
		long sum = 0;
		for (; ind > 0; ind -= ind & -ind) { sum += bit[ind]; }
		return sum;
	}
}
// EndCodeSnip