結果

問題 No.1002 Twotone
ユーザー Yu_212Yu_212
提出日時 2024-02-22 23:56:25
言語 C++23
(gcc 12.3.0 + boost 1.83.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
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 1 ms
5,248 KB
testcase_01 AC 2 ms
5,248 KB
testcase_02 AC 1 ms
5,248 KB
testcase_03 AC 1,048 ms
245,760 KB
testcase_04 AC 1,438 ms
321,860 KB
testcase_05 AC 1,471 ms
331,276 KB
testcase_06 AC 1 ms
6,816 KB
testcase_07 AC 667 ms
194,048 KB
testcase_08 AC 1,121 ms
321,612 KB
testcase_09 AC 1,211 ms
321,536 KB
testcase_10 AC 2 ms
6,820 KB
testcase_11 AC 977 ms
243,340 KB
testcase_12 AC 1,336 ms
316,624 KB
testcase_13 AC 1,311 ms
316,452 KB
testcase_14 AC 2 ms
6,816 KB
testcase_15 AC 680 ms
174,604 KB
testcase_16 AC 1,337 ms
316,408 KB
testcase_17 AC 1,299 ms
316,312 KB
testcase_18 AC 2 ms
6,816 KB
testcase_19 AC 1,115 ms
267,212 KB
testcase_20 AC 1,315 ms
313,472 KB
testcase_21 AC 1,329 ms
312,856 KB
testcase_22 AC 2 ms
6,816 KB
testcase_23 AC 824 ms
193,536 KB
testcase_24 AC 1,447 ms
321,096 KB
testcase_25 AC 1,425 ms
321,060 KB
testcase_26 AC 2 ms
6,816 KB
testcase_27 AC 809 ms
197,584 KB
testcase_28 AC 1,396 ms
301,284 KB
testcase_29 AC 1,355 ms
301,276 KB
testcase_30 AC 1 ms
6,816 KB
testcase_31 AC 1,314 ms
299,988 KB
testcase_32 AC 1,427 ms
301,204 KB
testcase_33 AC 1,362 ms
301,272 KB
testcase_34 AC 966 ms
315,312 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