RT,record
是这篇题解的 O(n+q) 做法
Code:
#include <iostream>
#include <vector>
// #include <algorithm>
using namespace std;
using ll = long long;
using ull = unsigned long long;
using pp = pair<int, int>;
#define pb push_back
const int mx = 1e6 + 2;
int n, q;
int color[mx];
int edge[mx];
void input() {
cin >> n >> q;
for (int i = 1; i <= n; ++i) cin >> color[i];
for (int i = 2; i <= n; ++i) cin >> edge[i];
}
struct bcj {
int fa[mx];
int sze[mx];
bcj() {
for (int i = 1; i < mx; ++i) fa[i] = i, sze[i] = 1;
}
int find(int x) {
if (fa[x] == x) return x;
return fa[x] = find(fa[x]);
}
void merge(int a, int b) {
int af = find(a);
int bf = find(b);
if (af != bf) fa[af] = bf, sze[bf] += sze[af];
}
};
ll samecolorcntans[mx];
const int md = 1e6 + 7;
int u[mx], ps[mx * 2];
int tag = 1, tot;
bcj samecolor = bcj();
bcj nwbcj = bcj();
struct colorpair {
pair<int, int> color;
int id;
};
vector<colorpair> unsamecolor[md + 2];
vector<vector<pair<int, int> > > unsamecoloredge;
ll unsamecolorans[mx];
pair<int, int> sortpair(pair<int, int> c) {
return {min(c.first, c.second), max(c.first, c.second)};
}
inline int gt_unsamecolorhash(pair<int, int> c) {
return ((1ll * mx * c.first) % md + c.second) % md;
}
colorpair *gt_unsamecolorid(pair<int, int> c) {
int cnt = 0;
for (auto &i : unsamecolor[gt_unsamecolorhash(c)]) {
if (c == i.color) return &i;
++cnt;
}
return nullptr;
}
int tryins_unsamecolorid(pair<int, int> c) {
auto tmp = gt_unsamecolorid(c);
if (tmp != nullptr) return (*tmp).id;
int hashd = gt_unsamecolorhash(c);
colorpair nw;
nw.color = c;
nw.id = unsamecoloredge.size();
unsamecolor[hashd].push_back(nw);
unsamecoloredge.push_back({});
return nw.id;
}
#define SQUARE(x) (1ll * (x) * (x))
void pre() {
// pre the same color's size in samecolorcntans[<color>]
for (int i = 2; i <= n; ++i)
if (color[i] == color[edge[i]]) samecolor.merge(i, edge[i]);
for (int i = 1; i <= n; ++i)
if (samecolor.find(i) == i)
samecolorcntans[color[i]] += SQUARE(samecolor.sze[i]);
// pre the 2 dots which connected by a edge
for (int i = 2; i <= n; ++i) {
if (color[i] != color[edge[i]]) {
int index = tryins_unsamecolorid(sortpair({color[i], color[edge[i]]}));
unsamecoloredge[index].push_back(
{samecolor.fa[i], samecolor.fa[edge[i]]});
}
}
int nwid = 0;
for (auto &nwedges : unsamecoloredge) {
tot = 0;
++tag;
for (auto [i, j] : nwedges) {
if (u[i] != tag) u[i] = tag, ps[tot++] = i;
if (u[j] != tag) u[j] = tag, ps[tot++] = j;
}
for (int i, ii = 0; ii < tot; ++ii) {
i = ps[ii];
nwbcj.fa[i] = i;
nwbcj.sze[i] = samecolor.sze[i];
unsamecolorans[nwid] -= SQUARE(samecolor.sze[i]);
}
for (auto [i, j] : nwedges) nwbcj.merge(i, j);
for (int i, ii = 0; ii < tot; ++ii) {
i = ps[ii];
if (nwbcj.find(i) == i) unsamecolorans[nwid] += SQUARE(nwbcj.sze[i]);
}
++nwid;
}
}
int x, y;
void reply_question() {
cin >> x >> y;
int ans = samecolorcntans[x] + samecolorcntans[y];
if (gt_unsamecolorid(sortpair({x, y})) != nullptr)
ans += unsamecolorans[(*gt_unsamecolorid(sortpair({x, y}))).id];
cout << ans << '\n';
}
int main() {
cin.tie(0);
cout.tie(0);
ios::sync_with_stdio(false);
input();
pre();
while (q--) reply_question();
}