[蓝桥杯 2024 省 A] 零食采购

[蓝桥杯 2024 省 A] 零食采购

题目描述

小蓝准备去星际旅行，出发前想在本星系采购一些零食，星系内有 n颗星球，由 n−1条航路连接为连通图，第 i颗星球卖第 ci种零食特产。小蓝想出了 q个采购方案，第 i个方案的起点为星球 si ，终点为星球 ti，对于每种采购方案，小蓝将从起点走最短的航路到终点，并且可以购买所有经过的星球上的零食（包括起点终点），请计算每种采购方案最多能买多少种不同的零食。

输入格式

输入的第一行包含两个正整数 n，q用一个空格分隔。
第二行包含 nn 个整数 c1,c2,⋯ ,cn相邻整数之间使用一个空格分隔。
接下来 n−1行，第 i行包含两个整数 ui,vi用一个空格分隔，表示一条 航路将星球 ui与 vi相连。
接下来 q行，第 i行包含两个整数 si,ti用一个空格分隔，表示一个采购方案。

输出格式

输出 q行，每行包含一个整数，依次表示每个采购方案的答案。

对于所有评测用例，1≤n,q≤105，1≤ci≤20，1≤ui,vi≤n，1≤si,ti≤n1≤n,q≤105，1≤ci​≤20，1≤ui​,vi​≤n，1≤si​,ti​≤n。

解析

首先将一个顶点作为树根，形成一个树。对于每一个查询s,t,如果能知道s和t的最近公共祖先fa,则只需计算s到fa购买了哪些零食，t到fa购买了哪些零食，求和即可。

因此我们需要LCA算法求出s,t的最近公共祖先，方便每次查询，时间复杂度为O(nlogn)。

详细代码见 最近公共祖先LCA算法（倍增算法）

其次我们需要计算s到fa购买了哪些零食，t到fa购买了哪些零食，进行求和。

可以用一个二维数组a[i][j]表示顶点i到根节点的路径上零食j有多少个（包括顶点i）

因此s到t的最短路径上，零食ci的个数为

a[s][i] + a[t][i] - 2 * a[fa][i] + (i == c[fa]);

构造二维数组的过程只需要在dfs的过程中，将父节点的a[fa]赋给子节点，并加上子节点的零食即可。

for (int i = 1; i <= 22; i++) a[v][i] = a[fa][i];
a[v][c[v]]++;

详细代码

#include <iostream>
#include <cstring>
#include <algorithm>

using namespace std;
const int maxsize = 1e5 + 10;
struct edge {
	int to, next;
}e[2 * maxsize];
int cnt = 0;
int head[maxsize];
int dep[maxsize];
int f[maxsize][21];
int n, q;
int c[maxsize];
int a[maxsize][23];
void add(int u, int v) {
	e[cnt].to = v;
	e[cnt].next = head[u];
	head[u] = cnt++;
}
void dfs(int v,int fa) {
	f[v][0] = fa; dep[v] = dep[fa] + 1;
	for (int i = 1; i <= 22; i++) a[v][i] = a[fa][i];
	a[v][c[v]]++;
	for (int k = 1; (1 << k) <= dep[v] - 1; k++) {
		f[v][k] = f[f[v][k - 1]][k - 1];
	}
	for (int i = head[v]; i != -1; i = e[i].next) {
		int u = e[i].to;
		if (u != fa) dfs(u, v);
	}
}
int lca(int x, int y) {
	if (dep[x] > dep[y]) swap(x, y);
	for (int k = 20; k >= 0; k--) {
		if (dep[x] + (1 << k) <= dep[y]) y = f[y][k];
	}
	
	if (x == y) return x;
	for (int k = 20; k >= 0; k--) {
		if (f[y][k] == f[x][k]) continue;
		y = f[y][k]; x = f[x][k];
	}
	return f[x][0];
}
int main()
{
	scanf("%d %d", &n, &q);
	for (int v = 1; v <= n; v++) scanf("%d", &c[v]);
	memset(head, -1, sizeof(head)); memset(a, 0, sizeof(a)); memset(f, 0, sizeof(f)); memset(dep, 0, sizeof(dep));
	int u, v;
	for (int i = 1; i < n; i++) {
		scanf("%d %d", &u, &v);
		add(u, v); add(v, u);
	}
	dfs(1, 0);
	while (q--) {
		scanf("%d %d", &u, &v);
		int fa = lca(u,v);
		int ans = 0;
		for (int item = 1; item <= 22; item++) {
			int val = a[u][item] + a[v][item] - 2 * a[fa][item] + (item == c[fa]);
			if (val > 0) ans++;
		}
		printf("%d\n", ans);
	}
	return 0;
}