結果

問題 No.1002 Twotone
ユーザー Yu_212
提出日時 2024-02-22 23:56:25
言語 C++23
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 1,471 ms / 5,000 ms
コード長 4,041 bytes
コンパイル時間 3,156 ms
コンパイル使用メモリ 279,124 KB
実行使用メモリ 331,276 KB
最終ジャッジ日時 2024-09-29 04:31:53
合計ジャッジ時間 34,950 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 33
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <bits/stdc++.h>
using namespace std;
using ll = long long;


using pii = pair<int, int>;
using graph = vector<vector<pii>>;
template <class F>
void one_third_centroid_decomposition(graph &g, F f) {
  int n = g.size();
  assert(n >= 1);
  vector<int> size(n), xy(n), count(n), sort(n);
  using buf_node = pair<graph, graph>;
  vector<unique_ptr<buf_node>> buf;
  auto get_buf = [&](int depth) -> buf_node& {
    while (depth >= buf.size())
      buf.push_back(make_unique<buf_node>(graph(n), graph(n)));
    return *buf[depth];
  };
  auto dc = [&](auto &dc, int n, int v, graph &g, int depth) -> void {
    if (n <= 2) return;
    int s = -1;
    auto sz = [&](auto &sz, int v, int p) -> void {
      size[v] = 1;
      int mx = 0;
      for (auto ee : g[v]) {
        int e = ee.first;
        if (e != p) {
          sz(sz, e, v);
          size[v] += size[e];
          mx = max(mx, size[e]);
        }
      }
      mx = max(mx, n - size[v]);
      if (mx * 2 <= n)
        s = v;
    };
    sz(sz, v, -1);
    int maxsz = 0;
    for (auto ee : g[s]) {
      int e = ee.first;
      if (size[e] > size[s])
        size[e] = n - size[s];
      count[size[e]]++;
      maxsz = max(maxsz, size[e] + 1);
    }
    for (int i = 1; i < maxsz; i++)
      count[i] += count[i - 1];
    for (auto ee : g[s]) {
      int e = ee.first;
      sort[--count[size[e]]] = e;
    }
    fill(count.begin(), count.begin() + maxsz, 0);
    int xs = 0, ys = 0;
    for (int i = g[s].size(); i-- > 0;) {
      if (xs < ys)
        xy[sort[i]] = 0, xs += size[sort[i]];
      else
        xy[sort[i]] = 1, ys += size[sort[i]];
    }
    auto &[x, y] = get_buf(depth);
    x[s].clear();
    y[s].clear();
    auto build = [&](auto &build, int v, int p, auto &t) -> void {
      t[v] = g[v];
      for (auto ee : g[v]) {
        int e = ee.first;
        if (e != p) build(build, e, v, t);
      }
    };
    for (auto ee : g[s]) {
      int e = ee.first;
      auto &t = xy[e] ? y : x;
      t[s].push_back(ee);
      build(build, e, s, t);
    }
    f(s, x, y);
    dc(dc, xs + 1, s, x, depth + 1);
    dc(dc, ys + 1, s, y, depth + 1);
  };
  dc(dc, n, 0, g, 0);
}

int main() {
  cin.tie(0)->sync_with_stdio(false);
  int n, k;
  cin >> n >> k;
  using pii = pair<int, int>;
  vector<vector<pii>> edges(n);
  for (int i = 0; i < n - 1; i++) {
    int u, v, c;
    cin >> u >> v >> c;
    u--, v--, c--;
    edges[u].emplace_back(v, c);
    edges[v].emplace_back(u, c);
  }
  ll ans = 0;
  one_third_centroid_decomposition(edges, [&](int ce, vector<vector<pii>> &x, vector<vector<pii>> &y) {
    int cnt0 = 0; // *, -
    int cnt1 = 0; // *, *
    map<pii, int> cnt2; // x, y
    map<int, int> cnt3; // x, *
    map<int, int> cnt4; // x, -
    auto f1 = [&](auto self, int node, int parent, int c1, int c2) -> void {
      if (c1 == -1) {
      } else if (c2 == -1) {
        cnt0++;
        cnt4[c1]++;
      } else {
        cnt1++;
        cnt2[make_pair(min(c1, c2), max(c1, c2))]++;
        cnt3[c1]++;
        cnt3[c2]++;
      }
      for (auto [child, c] : x[node]) {
        if (child == parent) continue;
        if (c1 == -1) {
          self(self, child, node, c, -1);
        } else if (c1 == c) {
          self(self, child, node, c1, c2);
        } else if (c2 == c || c2 == -1) {
          self(self, child, node, c1, c);
        }
      }
    };
    auto f2 = [&](auto self, int node, int parent, int c1, int c2) -> void {
      if (c1 == -1) {
      } else if (c2 == -1) {
        ans += cnt0 - cnt4[c1] + cnt3[c1];
      } else {
        ans += cnt2[make_pair(min(c1, c2), max(c1, c2))] + cnt4[c1] + cnt4[c2];
      }
      for (auto [child, c] : y[node]) {
        if (child == parent) continue;
        if (c1 == -1) {
          self(self, child, node, c, -1);
        } else if (c1 == c) {
          self(self, child, node, c1, c2);
        } else if (c2 == c || c2 == -1) {
          self(self, child, node, c1, c);
        }
      }
    };
    f1(f1, ce, -1, -1, -1);
    f2(f2, ce, -1, -1, -1);
  });
  cout << ans << endl;
}
0