結果

問題 No.1002 Twotone
ユーザー Yu_212Yu_212
提出日時 2024-02-22 23:56:25
言語 C++23
(gcc 12.3.0 + boost 1.83.0)
結果
AC  
実行時間 1,837 ms / 5,000 ms
コード長 4,041 bytes
コンパイル時間 5,831 ms
コンパイル使用メモリ 278,916 KB
実行使用メモリ 331,240 KB
最終ジャッジ日時 2024-02-22 23:57:12
合計ジャッジ時間 42,576 ms
ジャッジサーバーID
（参考情報） 		judge14 / judge13
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 2 ms
6,676 KB
testcase_01 AC 2 ms
6,676 KB
testcase_02 AC 2 ms
6,676 KB
testcase_03 AC 1,187 ms
245,884 KB
testcase_04 AC 1,664 ms
321,780 KB
testcase_05 AC 1,630 ms
331,240 KB
testcase_06 AC 2 ms
6,676 KB
testcase_07 AC 820 ms
194,240 KB
testcase_08 AC 1,345 ms
321,724 KB
testcase_09 AC 1,337 ms
321,640 KB
testcase_10 AC 2 ms
6,676 KB
testcase_11 AC 1,189 ms
243,484 KB
testcase_12 AC 1,612 ms
316,432 KB
testcase_13 AC 1,641 ms
316,432 KB
testcase_14 AC 2 ms
6,676 KB
testcase_15 AC 821 ms
174,648 KB
testcase_16 AC 1,645 ms
316,428 KB
testcase_17 AC 1,571 ms
316,428 KB
testcase_18 AC 3 ms
6,676 KB
testcase_19 AC 1,356 ms
267,196 KB
testcase_20 AC 1,593 ms
313,444 KB
testcase_21 AC 1,628 ms
313,060 KB
testcase_22 AC 2 ms
6,676 KB
testcase_23 AC 1,003 ms
193,548 KB
testcase_24 AC 1,837 ms
321,164 KB
testcase_25 AC 1,825 ms
321,164 KB
testcase_26 AC 2 ms
6,676 KB
testcase_27 AC 997 ms
197,636 KB
testcase_28 AC 1,815 ms
301,416 KB
testcase_29 AC 1,672 ms
301,416 KB
testcase_30 AC 2 ms
6,676 KB
testcase_31 AC 1,619 ms
299,944 KB
testcase_32 AC 1,778 ms
301,408 KB
testcase_33 AC 1,631 ms
301,408 KB
testcase_34 AC 1,128 ms
315,432 KB
権限があれば一括ダウンロードができます

ソースコード

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