test/atcoder/abc/abc214/D.test.cpp

View this file on GitHub
Last update: 2023-07-08 14:12:19+09:00
Problem: https://atcoder.jp/contests/abc214/tasks/abc214_d

Depends on

Code

#define PROBLEM "https://atcoder.jp/contests/abc214/tasks/abc214_d"

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

#include "data-structure/union-find-undo.hpp"
using namespace matumoto;


// {{{ Templates

// clang-format off

// Macros
#define over_load_(_1,_2,_3,_4,NAME,...) NAME
#define rep(...) over_load_(__VA_ARGS__, rep4, rep3, rep2)(__VA_ARGS__)
#define rep2(i, r) for ( int i = 0; i < static_cast<int>(r); (i) += 1)
#define rep3(i, l, r) for ( int i = static_cast<int>(l); i < static_cast<int>(r); (i) += 1)
#define rep4(i, l, r, stride) for ( int i = static_cast<int>(l); i < static_cast<int>(r); (i) += (stride))
#define rrep(...) over_load_(__VA_ARGS__, rrep4, rrep3, rrep2)(__VA_ARGS__)
#define rrep2(i, r) for ( int i = static_cast<int>(r) - 1; i >= 0; (i) -= 1)
#define rrep3(i, l, r) for ( int i = static_cast<int>(r) - 1; i >= static_cast<int>(l); (i) -= 1)
#define rrep4(i, l, r, stride) for ( int i = static_cast<int>(r) - 1; i >= static_cast<int>(l); (i) -= (stride))
#define len(x) (static_cast<int>((x).size()))
#define whole(f, x, ...) ([&](decltype((x)) container) { return (f)( begin(container), end(container), ## __VA_ARGS__); })(x)
#define rwhole(f, x, ...) ([&](decltype((x)) container) { return (f)( rbegin(container), rend(container), ## __VA_ARGS__); })(x)
#define debug(...) debug_function(#__VA_ARGS__, __VA_ARGS__)

// Operators
template <typename T1, typename T2> ostream &operator<<(ostream &os, const pair<T1, T2> &p) { os << "(" << p.first << "," << p.second << ")"; return os; }
template <typename T1, typename T2> ostream &operator<<(ostream &os, const map<T1, T2> &v) { bool is_first = true; for (auto x: v) { os << (is_first ? "" : " ") << x; is_first = false; } return os; }
template <typename T> ostream &operator<<(ostream &os, queue<T> v) { bool is_first = true; while (!v.empty()) { os << (is_first?"":" ")<<v.front(); v.pop(); is_first = false; } return os; }
template <typename T> ostream &operator<<(ostream &os, stack<T> v) { bool is_first = true; while (!v.empty()) { os << (is_first?"":" ") << v.top(); v.pop(); is_first=false; } return os; }
template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) { rep (i, len(v)) os << v[i] << (i == len(v) - 1 ? "" : " "); return os; }
template <typename T> ostream &operator<<(ostream &os, const vector<vector<T>> &v) { for (const auto &vec: v) { os << vec << '\n'; } return os; }
template <typename T> ostream &operator<<(ostream &os, const deque<T> &v) { rep (i, len(v)) os << v[i] << (i == len(v) - 1 ? "" : " "); return os; }
template <typename T> ostream &operator<<(ostream &os, const set<T> &v) { bool is_first = true; for (T x: v) { os << (is_first ? "" : " ") << x; is_first = false; } return os; }
template <typename T> istream &operator>>(istream &is, vector<T> &v) { for (T &in: v) { is >> in; } return is; }

// For debug macro
int find_comma_not_bracketed(string_view s){ stack<char> bs; string lbs = "({[", rbs = ")}]"; for (size_t i = 0; i < s.size(); i++) { if (lbs.find(s[i]) != string::npos) bs.push(s[i]); if (rbs.find(s[i]) != string::npos and !bs.empty()) bs.pop(); if (s[i] == ',' and bs.empty()) return i; } return s.size(); }
template <typename T, typename... Ts> void debug_function(string_view name, const T &a, Ts &&...rest) { int end = find_comma_not_bracketed(name); cerr << name.substr(0, end) << ":" << a; if constexpr (sizeof...(rest) == 0) { cerr << '\n'; } else { cerr << ' '; debug_function(name.substr(name.find_first_not_of(' ', end + 1)), forward<Ts>(rest)...); } }

// Functions
template <typename T> vector<T> make_vector(size_t a, T b) { return vector<T>(a, b); }
template <typename... Ts> auto make_vector(size_t a, Ts... ts) { return vector<decltype(make_vector(ts...))>(a, make_vector(ts...)); }
template <typename T1, typename T2> inline bool chmax(T1 &a, T2 b) { return a < b and (a = b, true); }
template <typename T1, typename T2> inline bool chmin(T1 &a, T2 b) { return a > b and (a = b, true); }

// Structs
struct IoSetup { IoSetup(int x = 15) { cin.tie(nullptr); ios::sync_with_stdio(false); cout << fixed << setprecision(x); cerr << fixed << setprecision(x); } } iosetup;

// Type aliases
using ull = unsigned long long;
using ll = long long;
using pll = pair<ll, ll>;
using pii = pair<int, int>;

// Literals
constexpr ll INF64 = INT64_MAX / 2;
constexpr int INF32 = INT32_MAX / 2;
constexpr int dy[] = { 0, 1, -1, 0, -1, 1, -1, 1 };
constexpr int dx[] = { 1, 0, 0, -1, -1, -1, 1, 1 };
constexpr int mod998244353 = 998244353;
constexpr int mod1000000007 = static_cast<int>(1e9) + 7;
constexpr char newl = '\n';

// clang-format on

// }}} Templates


int main() {
  int n;
  cin >> n;

  // <w, u, v>
  vector<tuple<int, int, int>> edges(n - 1);

  UnionFindUndo uf(n);

  rep(i, n - 1) {
    int u, v, w;
    cin >> u >> v >> w;

    u--, v--;

    edges[i] = tuple(w, u, v);
  }

  whole(sort, edges);

  rep(i, n - 1) {
    auto [w, u, v] = edges[i];
    uf.merge(u, v);
  }


  ll ans = 0;

  rrep(i, n - 1) {
    auto [w, u, v] = edges[i];

    uf.undo();
    ans += ll(uf.size(u)) * ll(uf.size(v)) * ll(w);
  }

  cout << ans << newl;
}

#line 1 "test/atcoder/abc/abc214/D.test.cpp"
#define PROBLEM "https://atcoder.jp/contests/abc214/tasks/abc214_d"

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

#line 2 "data-structure/union-find-undo.hpp"

#line 2 "data-structure/base.hpp"

namespace matumoto {
  using namespace std;
}
#line 4 "data-structure/union-find-undo.hpp"

#line 9 "data-structure/union-find-undo.hpp"

namespace matumoto {
  // 経路圧縮なし
  class UnionFindUndo {
  private:
    // number of vertices
    int n_;

    vector<int> sizes_, parents_;

    // <index, parent(index), size(index)>
    stack<tuple<int, int, int>> history_;

  public:
    UnionFindUndo(int N): sizes_(N, 1), parents_(N) {
      iota(parents_.begin(), parents_.end(), 0);
    }

    // 根（そのグループの識別番号）
    int root(int x) {
      if (x == parents_[x])
        return x;
      return root(parents_[x]);
    }

    int size(int x) {
      return sizes_[root(x)];
    }

    bool same(int x, int y) {
      return root(x) == root(y);
    }

    // size(x) < size(y): -1
    // not merged:         0
    // size(x) > size(y): +1
    int merge(int x, int y) {
      x = root(x);
      y = root(y);
      if (x == y)
        return 0;

      int result = +1;
      if (sizes_[x] < sizes_[y]) {
        swap(x, y);
        result = -1;
      }

      history_.emplace(tuple(x, parents_[x], sizes_[x]));
      history_.emplace(tuple(y, parents_[y], sizes_[y]));
      sizes_[x] += sizes_[y];
      parents_[y] = x;
      return result;
    }

    bool undo() {
      if (history_.empty())
        return false;

      auto [x, x_parent, x_size] = history_.top();
      history_.pop();
      auto [y, y_parent, y_size] = history_.top();
      history_.pop();

      parents_[x] = x_parent;
      sizes_[x] = x_size;
      parents_[y] = y_parent;
      sizes_[y] = y_size;
      return true;
    }

    void clear_history() {
      while (not history_.empty()) {
        history_.pop();
      }
    }

    void all_undo() {
      while (undo()) {}
    }

    // Θ(N log(N))
    vector<vector<int>> groups() {
      vector<vector<int>> grps(n_);
      for (int i = 0; i < n_; i++) {
        grps[root(i)].emplace_back(i);
      }

      vector<vector<int>> res;
      for (int i = 0; i < n_; i++) {
        if (grps[i].empty())
          continue;
        res.emplace_back(grps[i]);
      }
      return res;
    }
  };
} // namespace matumoto
#line 7 "test/atcoder/abc/abc214/D.test.cpp"
using namespace matumoto;


// {{{ Templates

// clang-format off

// Macros
#define over_load_(_1,_2,_3,_4,NAME,...) NAME
#define rep(...) over_load_(__VA_ARGS__, rep4, rep3, rep2)(__VA_ARGS__)
#define rep2(i, r) for ( int i = 0; i < static_cast<int>(r); (i) += 1)
#define rep3(i, l, r) for ( int i = static_cast<int>(l); i < static_cast<int>(r); (i) += 1)
#define rep4(i, l, r, stride) for ( int i = static_cast<int>(l); i < static_cast<int>(r); (i) += (stride))
#define rrep(...) over_load_(__VA_ARGS__, rrep4, rrep3, rrep2)(__VA_ARGS__)
#define rrep2(i, r) for ( int i = static_cast<int>(r) - 1; i >= 0; (i) -= 1)
#define rrep3(i, l, r) for ( int i = static_cast<int>(r) - 1; i >= static_cast<int>(l); (i) -= 1)
#define rrep4(i, l, r, stride) for ( int i = static_cast<int>(r) - 1; i >= static_cast<int>(l); (i) -= (stride))
#define len(x) (static_cast<int>((x).size()))
#define whole(f, x, ...) ([&](decltype((x)) container) { return (f)( begin(container), end(container), ## __VA_ARGS__); })(x)
#define rwhole(f, x, ...) ([&](decltype((x)) container) { return (f)( rbegin(container), rend(container), ## __VA_ARGS__); })(x)
#define debug(...) debug_function(#__VA_ARGS__, __VA_ARGS__)

// Operators
template <typename T1, typename T2> ostream &operator<<(ostream &os, const pair<T1, T2> &p) { os << "(" << p.first << "," << p.second << ")"; return os; }
template <typename T1, typename T2> ostream &operator<<(ostream &os, const map<T1, T2> &v) { bool is_first = true; for (auto x: v) { os << (is_first ? "" : " ") << x; is_first = false; } return os; }
template <typename T> ostream &operator<<(ostream &os, queue<T> v) { bool is_first = true; while (!v.empty()) { os << (is_first?"":" ")<<v.front(); v.pop(); is_first = false; } return os; }
template <typename T> ostream &operator<<(ostream &os, stack<T> v) { bool is_first = true; while (!v.empty()) { os << (is_first?"":" ") << v.top(); v.pop(); is_first=false; } return os; }
template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) { rep (i, len(v)) os << v[i] << (i == len(v) - 1 ? "" : " "); return os; }
template <typename T> ostream &operator<<(ostream &os, const vector<vector<T>> &v) { for (const auto &vec: v) { os << vec << '\n'; } return os; }
template <typename T> ostream &operator<<(ostream &os, const deque<T> &v) { rep (i, len(v)) os << v[i] << (i == len(v) - 1 ? "" : " "); return os; }
template <typename T> ostream &operator<<(ostream &os, const set<T> &v) { bool is_first = true; for (T x: v) { os << (is_first ? "" : " ") << x; is_first = false; } return os; }
template <typename T> istream &operator>>(istream &is, vector<T> &v) { for (T &in: v) { is >> in; } return is; }

// For debug macro
int find_comma_not_bracketed(string_view s){ stack<char> bs; string lbs = "({[", rbs = ")}]"; for (size_t i = 0; i < s.size(); i++) { if (lbs.find(s[i]) != string::npos) bs.push(s[i]); if (rbs.find(s[i]) != string::npos and !bs.empty()) bs.pop(); if (s[i] == ',' and bs.empty()) return i; } return s.size(); }
template <typename T, typename... Ts> void debug_function(string_view name, const T &a, Ts &&...rest) { int end = find_comma_not_bracketed(name); cerr << name.substr(0, end) << ":" << a; if constexpr (sizeof...(rest) == 0) { cerr << '\n'; } else { cerr << ' '; debug_function(name.substr(name.find_first_not_of(' ', end + 1)), forward<Ts>(rest)...); } }

// Functions
template <typename T> vector<T> make_vector(size_t a, T b) { return vector<T>(a, b); }
template <typename... Ts> auto make_vector(size_t a, Ts... ts) { return vector<decltype(make_vector(ts...))>(a, make_vector(ts...)); }
template <typename T1, typename T2> inline bool chmax(T1 &a, T2 b) { return a < b and (a = b, true); }
template <typename T1, typename T2> inline bool chmin(T1 &a, T2 b) { return a > b and (a = b, true); }

// Structs
struct IoSetup { IoSetup(int x = 15) { cin.tie(nullptr); ios::sync_with_stdio(false); cout << fixed << setprecision(x); cerr << fixed << setprecision(x); } } iosetup;

// Type aliases
using ull = unsigned long long;
using ll = long long;
using pll = pair<ll, ll>;
using pii = pair<int, int>;

// Literals
constexpr ll INF64 = INT64_MAX / 2;
constexpr int INF32 = INT32_MAX / 2;
constexpr int dy[] = { 0, 1, -1, 0, -1, 1, -1, 1 };
constexpr int dx[] = { 1, 0, 0, -1, -1, -1, 1, 1 };
constexpr int mod998244353 = 998244353;
constexpr int mod1000000007 = static_cast<int>(1e9) + 7;
constexpr char newl = '\n';

// clang-format on

// }}} Templates


int main() {
  int n;
  cin >> n;

  // <w, u, v>
  vector<tuple<int, int, int>> edges(n - 1);

  UnionFindUndo uf(n);

  rep(i, n - 1) {
    int u, v, w;
    cin >> u >> v >> w;

    u--, v--;

    edges[i] = tuple(w, u, v);
  }

  whole(sort, edges);

  rep(i, n - 1) {
    auto [w, u, v] = edges[i];
    uf.merge(u, v);
  }


  ll ans = 0;

  rrep(i, n - 1) {
    auto [w, u, v] = edges[i];

    uf.undo();
    ans += ll(uf.size(u)) * ll(uf.size(v)) * ll(w);
  }

  cout << ans << newl;
}