P5838 [USACO19DEC] Milk Visits G

P5838 [USACO19DEC] Milk Visits G

Luogu P5838

Solution

提供一种奇特的 \(\mathcal O(\dfrac{n\sqrt n\log n}{\omega})\) 的做法。

树链剖分转化成序列问题。然后变成了询问一个区间 \(l,r\) 是否存在一个颜色 \(k\)，显然可以 \(\mathcal O(n)\) 预处理 \(\mathcal O(\sqrt n)\) 判定，这样复杂度是 \(\mathcal O(n\sqrt n\log n)\) 的，空间复杂度 \(\mathcal O(n\sqrt n)\)，很卡，没敢写。

注意到并不需要数颜色出现了多少次，因此考虑使用 bitset 维护一个颜色是否出现。

对于整块，开 \(\mathcal O(n)\) 个 bitset，其中第 \(i\) 个 bitset 的第 \(j\) 位表示颜色 \(i\) 是否在第 \(j\) 块出现。那么对于整块，可以使用 bitset 的位运算在 \(\mathcal O(\dfrac{\sqrt n}{\omega})\) 的时间内判定。

对于散块，考虑对每个块中的每一个颜色开一个长度为块长的 bitset ，第 \(i\) 位表示这个颜色在这个块中从左起的第 \(i\) 个位置出现了，那么询问也容易使用 bitset 的位运算做到 \(\mathcal O(\dfrac{\sqrt n}{\omega})\)。乍一看这部分的 bitset 大小是 \(\mathcal O(\dfrac{n^2}{\omega})\) 的，但是因为一个块中的颜色数量是 \(\mathcal O(\sqrt n)\) 的，所以可以将这个块内的颜色重新映射到 \([1,\sqrt n]\) 的范围（类似于离散化），然后空间复杂度就降为了 \(\mathcal O(\dfrac{n\sqrt n}{\omega})\)。颜色在每个块内的离散化可以使用 Hash 表做到线性处理。

总时间复杂度 \(\mathcal O(\dfrac{n\sqrt n\log n}{\omega})\)，空间复杂度 \(\mathcal O(\dfrac{n\sqrt n}{\omega})\)。

#include<bits/stdc++.h>
#define For(i, a, b) for (int i = (a); i <= (b); ++i)
#define Rof(i, a, b) for (int i = (a); i >= (b); --i)
using namespace std;
using i64 = long long;
constexpr int _N = 1e5 + 5, _B = 400 + 5;
bool ST;
int n, m;
vector<int> e[_N];
int dep[_N], fa[_N], id[_N], son[_N];
int top[_N], siz[_N], dfn;
int a[_N], b[_N];
void Dfs1(int x, int F) {
	dep[x] = dep[F] + 1, fa[x] = F, siz[x] = 1;
	int maxson = -1;
	for (int v : e[x]) if (v != F) {
		Dfs1(v, x);
		if (siz[v] > maxson) maxson = siz[v], son[x] = v;
		siz[x] += siz[v];
	}
}
void Dfs2(int x, int topf) {
	id[x] = ++dfn, top[x] = topf, b[dfn] = a[x];
	if (son[x]) Dfs2(son[x], topf);
	for (int v : e[x]) if (v != son[x] && v != fa[x])
		Dfs2(v, v);
}
int pos[_N], bl[_N], br[_N], len, cnt;
bitset<_B> ct[_N], tl[_B], bct[_B][_B];
int tot[_B];
#include<ext/pb_ds/assoc_container.hpp>
__gnu_pbds::gp_hash_table<int, int> H[_B];
bool ED;
void InitBlock() {
	len = sqrt(n), cnt = (n - 1) / len + 1;
	For(i, 1, cnt) bl[i] = br[i - 1] + 1, br[i] = min(i * len, n);
	For(i, 1, cnt) tl[i] = tl[i - 1], tl[i].set(i);
	For(i, 1, cnt) {
		For(j, bl[i], br[i]) {
			pos[j] = i;
			ct[b[j]].set(i);
			if (H[i].find(b[j]) == H[i].end())
				H[i][b[j]] = ++tot[i];
			int nc = H[i][b[j]];
			bct[i][nc].set(j - bl[i] + 1);
		}
	}
}
bool AskB(int id, int l, int r, int k) {
	if (H[id].find(k) == H[id].end()) return 0;
	int nc = H[id][k];
	return ((bct[id][nc] >> (l - bl[id])) & tl[r - l + 1]).any();
}
bool Query(int l, int r, int k) {
	if (pos[l] == pos[r])
		return AskB(pos[l], l, r, k);
	else
		return AskB(pos[l], l, br[pos[l]], k)
			|| AskB(pos[r], bl[pos[r]], r, k)
			|| ((ct[k] >> pos[l]) & tl[pos[r] - pos[l] - 1]).any();
}
bool QueryPath(int x, int y, int k) {
	bool res = 0;
	while (top[x] != top[y]) {
		if (dep[top[x]] < dep[top[y]]) swap(x, y);
		res = Query(id[top[x]], id[x], k);
		if (res) return res;
		x = fa[top[x]];
	}
	if (dep[x] > dep[y]) swap(x, y);
	return Query(id[x], id[y], k);
}
signed main() {
	ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
	cin >> n >> m;
	For(i, 1, n) cin >> a[i];
	For(i, 2, n) {
		int x, y;
		cin >> x >> y;
		e[x].push_back(y);
		e[y].push_back(x);
	}
	Dfs1(1, 0), Dfs2(1, 1);
	InitBlock();
	For(i, 0, m - 1) {
		int x, y, k;
		cin >> x >> y >> k;
		cout << QueryPath(x, y, k);
	}
}