結果

問題 No.1002 Twotone
ユーザー 0w10w1
提出日時 2021-01-02 16:51:50
言語 C++17
(gcc 12.3.0 + boost 1.83.0)
結果
AC  
実行時間 3,149 ms / 5,000 ms
コード長 4,510 bytes
コンパイル時間 3,140 ms
コンパイル使用メモリ 245,772 KB
実行使用メモリ 126,024 KB
最終ジャッジ日時 2024-10-12 06:00:06
合計ジャッジ時間 38,418 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge1
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 2 ms
5,248 KB
testcase_01 AC 33 ms
35,956 KB
testcase_02 AC 29 ms
32,088 KB
testcase_03 AC 692 ms
85,912 KB
testcase_04 AC 971 ms
90,584 KB
testcase_05 AC 1,019 ms
120,832 KB
testcase_06 AC 23 ms
25,628 KB
testcase_07 AC 650 ms
46,080 KB
testcase_08 AC 1,198 ms
95,644 KB
testcase_09 AC 1,225 ms
89,544 KB
testcase_10 AC 10 ms
10,496 KB
testcase_11 AC 1,498 ms
89,988 KB
testcase_12 AC 2,037 ms
73,996 KB
testcase_13 AC 2,062 ms
102,016 KB
testcase_14 AC 13 ms
13,272 KB
testcase_15 AC 974 ms
54,452 KB
testcase_16 AC 2,169 ms
82,176 KB
testcase_17 AC 2,164 ms
76,928 KB
testcase_18 AC 40 ms
43,136 KB
testcase_19 AC 1,585 ms
103,808 KB
testcase_20 AC 1,878 ms
88,424 KB
testcase_21 AC 1,838 ms
87,996 KB
testcase_22 AC 12 ms
12,888 KB
testcase_23 AC 1,598 ms
66,048 KB
testcase_24 AC 3,149 ms
108,020 KB
testcase_25 AC 3,118 ms
107,656 KB
testcase_26 AC 8 ms
7,808 KB
testcase_27 AC 205 ms
97,388 KB
testcase_28 AC 565 ms
95,476 KB
testcase_29 AC 298 ms
108,428 KB
testcase_30 AC 10 ms
9,984 KB
testcase_31 AC 280 ms
107,296 KB
testcase_32 AC 564 ms
93,116 KB
testcase_33 AC 321 ms
118,796 KB
testcase_34 AC 633 ms
126,024 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

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

/*^ vector extensions  */

template<typename T>
vector<T> concat(vector<T> a, vector<T> b) {
  a.insert(a.end(), b.begin(), b.end());
  return a;
}

template<typename T, size_t sz>
struct _Matrix_type { typedef vector<typename _Matrix_type<T, sz - 1>::type> type; };
template<typename T>
struct _Matrix_type<T, 1> { typedef T type; };
template<typename T>
struct _Matrix {
  static auto build(size_t s) { return vector<T>(s); }
  template<typename ...Args>
  static auto build(size_t f, Args... args) {
    return vector<typename _Matrix_type<T, 1 + sizeof...(args)>::type>(f, _Matrix<T>::build(args...));
  }
};
template<typename T, typename... Args>
auto buildMatrix(Args... args) { return _Matrix<T>::build(args...); }

/*  vector extensions $*/

/*^ generic definitions  */

template<typename F>
struct _RecurFun : F {
  _RecurFun(F&& f) : F(forward<F>(f)) {}
  template<typename... Args>
  decltype(auto) operator()(Args&&... args) const { return F::operator()(*this, forward<Args>(args)...); }
};

template<typename F>
decltype(auto) RecurFun(F&& f) { return _RecurFun<F> { forward<F>(f) }; }

/*  generic definitions $*/

/*^ pbds  */

#if __has_include(<ext/pb_ds/assoc_container.hpp>)
  #include <ext/pb_ds/assoc_container.hpp>
  using namespace __gnu_pbds;

  const int RandomIntSeed = chrono::high_resolution_clock::now().time_since_epoch().count();
  struct HashInt { int operator()(int x) const { return x ^ RandomIntSeed; } };

  template<typename T, typename U, typename H = hash<T>>
  using Gmap = gp_hash_table<T, U, HashInt>;
  // using Gmap = gp_hash_table<T, U, H>;
#endif

/*  pbds $*/

struct Centroid {
  vector<vector<int>> graph;
  vector<int> size, dead;
  Centroid(const vector<vector<int>> &g) : size(g.size()), dead(g.size()), graph(g) {}
  vector<int> findCentroids(int root) {
    vector<int> centroids;
    RecurFun([&](auto buildSize, int u, int p) -> void {
      size[u] = 1;
      for (int v : graph[u]) if (v != p && !dead[v]) buildSize(v, u), size[u] += size[v];
    })(root, -1);
    RecurFun([&](auto buildCentroids, int u, int p) -> void {
      int ok = (size[root] - size[u]) * 2 <= size[root];
      for (int v : graph[u]) if (v != p && !dead[v]) buildCentroids(v, u), ok &= size[v] * 2 <= size[root];
      if (ok) centroids.push_back(u);
    })(root, -1);
    assert(1 <= centroids.size() && centroids.size() <= 2);
    return centroids;
  }
};

int main() {
  ios::sync_with_stdio(false);

  int N, K; {
    cin >> N >> K;
  }

  vector<vector<int>> G(N); 
  vector<Gmap<int, int>> C(N); {
    for (int u = 0; u + 1 < N; ++u) {
      int U, V, Z; { 
        cin >> U >> V >> Z;
        --U, --V, --Z;
      }
      G[U].push_back(V);
      G[V].push_back(U);
      C[U][V] = Z;
      C[V][U] = Z;
    }
  }

  int64_t res = 0; {
    Centroid cd(G);

    queue<int> que({ cd.findCentroids(0)[0] });
    int all = 0; // single color
    vector<int> cnt(K), cntx(K);
    vector<Gmap<int, int>> cnt2(K);
    while (que.size()) {
      int u = que.front();
      que.pop();

      for (int v : G[u]) if (!cd.dead[v]) {
        RecurFun([&](auto count, int u, int p, int a, int b) -> void {
          for (int v : G[u]) if (v != p && !cd.dead[v]) {
            if (a == C[u][v] || b == C[u][v]) count(v, u, a, b);
            else if (a == -1) count(v, u, C[u][v], b);
          }
          if (a == -1) res += (all - cnt[b]) + cntx[b];
          else res += 1 + cnt[a] + cnt[b] + cnt2[a][b];
        })(v, u, -1, C[u][v]);

        RecurFun([&](auto update, int u, int p, int a, int b) -> void {
          for (int v : G[u]) if (v != p && !cd.dead[v]) {
            if (a == C[u][v] || b == C[u][v]) update(v, u, a, b);
            else if (a == -1) update(v, u, C[u][v], b);
          }
          if (a == -1) ++cnt[b], ++all;
          else ++cntx[a], ++cntx[b], ++cnt2[a][b], ++cnt2[b][a];
        })(v, u, -1, C[u][v]);
      }

      for (int v : G[u]) if (!cd.dead[v]) {
        RecurFun([&](auto update, int u, int p, int a, int b) -> void {
          for (int v : G[u]) if (v != p && !cd.dead[v]) {
            if (a == C[u][v] || b == C[u][v]) update(v, u, a, b);
            else if (a == -1) update(v, u, C[u][v], b);
          }
          if (a == -1) --cnt[b], --all;
          else --cntx[a], --cntx[b], --cnt2[a][b], --cnt2[b][a];
        })(v, u, -1, C[u][v]);
      }

      cd.dead[u] = 1;
      for (int v : G[u]) if (!cd.dead[v]) que.push(cd.findCentroids(v)[0]);
    }
  }

  cout << res << endl;
}
0