WA on #6 求助
查看原帖
WA on #6 求助
734630
justin_jiajia蒟蒻楼主2024/2/24 20:28

RT,record

是这篇题解的 O(n+q)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();
}
2024/2/24 20:28
加载中...