test/verify/yukicoder-1002.test.cpp

• View this file on GitHub
• Last update: 2024-05-01 23:44:47+09:00
• Problem: https://yukicoder.me/problems/no/1002

Depends on

• Graph Template(グラフテンプレート) (graph/graph-template.hpp)
• Centroid-Decomosition(重心分解) (graph/tree/centroid-decomposition.hpp)
• template/template.hpp

Code

// competitive-verifier: PROBLEM https://yukicoder.me/problems/no/1002

#include "../../template/template.hpp"

#include "../../graph/tree/centroid-decomposition.hpp"

int main() {
  int N, K;
  cin >> N >> K;
  CentroidDecomposition< int > g(N);
  g.read(N - 1, -1, true);
  int root = g.build();
  int64 ret = 0;
  vector< int > used(N);

  map< pair< int, int >, int > mp;
  int all;
  map< int, int > mp2;
  map< int, int > mp3;


  auto get_vec = MFP([&](auto get_vec, int idx, int par, int a, int b) -> void {
    if(b == -1) {
      ret += all - mp2[a];
      ret += mp3[a];
    } else {
      ret++;
      ret += mp2[a] + mp2[b];
      ret += mp[minmax(a, b)];
    }
    for(auto &e : g.g[idx]) {
      if(e.to == par) continue;
      if(used[e.to]) continue;
      if(a == e.cost) {
        get_vec(e.to, idx, e.cost, b);
      } else if(b == -1 || b == e.cost) {
        get_vec(e.to, idx, a, e.cost);
      }
    }
  });

  auto add_vec = MFP([&](auto add_vec, int idx, int par, int a, int b) -> void {
    if(b == -1) {
      mp2[a]++;
      all++;
    } else {
      mp[minmax(a, b)]++;
      mp3[a]++;
      mp3[b]++;
    }
    for(auto &e : g.g[idx]) {
      if(e.to == par) continue;
      if(used[e.to]) continue;
      if(a == e.cost) {
        add_vec(e.to, idx, e.cost, b);
      } else if(b == -1 || b == e.cost) {
        add_vec(e.to, idx, a, e.cost);
      }
    }
  });

  auto rec = MFP([&](auto rec, int idx) -> void {
    used[idx] = true;
    mp.clear();
    mp2.clear();
    mp3.clear();
    all = 0;
    for(auto &e : g.g[idx]) {
      if(used[e.to]) continue;
      get_vec(e.to, idx, e.cost, -1);
      add_vec(e.to, idx, e.cost, -1);
    }
    for(auto &to : g.tree.g[idx]) {
      rec(to);
    }
    used[idx] = false;
  });
  rec(root);
  cout << ret << endl;
}

#line 1 "test/verify/yukicoder-1002.test.cpp"
// competitive-verifier: PROBLEM https://yukicoder.me/problems/no/1002

#line 1 "template/template.hpp"
#include <bits/stdc++.h>

using namespace std;

using int64 = long long;

const int64 infll = (1LL << 62) - 1;
const int inf = (1 << 30) - 1;

struct IoSetup {
  IoSetup() {
    cin.tie(nullptr);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(10);
    cerr << fixed << setprecision(10);
  }
} iosetup;

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>
istream &operator>>(istream &is, pair<T1, T2> &p) {
  is >> p.first >> p.second;
  return is;
}

template <typename T>
ostream &operator<<(ostream &os, const vector<T> &v) {
  for (int i = 0; i < (int)v.size(); i++) {
    os << v[i] << (i + 1 != v.size() ? " " : "");
  }
  return os;
}

template <typename T>
istream &operator>>(istream &is, vector<T> &v) {
  for (T &in : v) is >> in;
  return is;
}

template <typename T1, typename T2>
inline bool chmax(T1 &a, T2 b) {
  return a < b && (a = b, true);
}

template <typename T1, typename T2>
inline bool chmin(T1 &a, T2 b) {
  return a > b && (a = b, true);
}

template <typename T = int64>
vector<T> make_v(size_t a) {
  return vector<T>(a);
}

template <typename T, typename... Ts>
auto make_v(size_t a, Ts... ts) {
  return vector<decltype(make_v<T>(ts...))>(a, make_v<T>(ts...));
}

template <typename T, typename V>
typename enable_if<is_class<T>::value == 0>::type fill_v(T &t, const V &v) {
  t = v;
}

template <typename T, typename V>
typename enable_if<is_class<T>::value != 0>::type fill_v(T &t, const V &v) {
  for (auto &e : t) fill_v(e, v);
}

template <typename F>
struct FixPoint : F {
  explicit FixPoint(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>
inline decltype(auto) MFP(F &&f) {
  return FixPoint<F>{forward<F>(f)};
}
#line 4 "test/verify/yukicoder-1002.test.cpp"

#line 2 "graph/tree/centroid-decomposition.hpp"

#line 2 "graph/graph-template.hpp"

/**
 * @brief Graph Template(グラフテンプレート)
 */
template <typename T = int>
struct Edge {
  int from, to;
  T cost;
  int idx;

  Edge() = default;

  Edge(int from, int to, T cost = 1, int idx = -1)
      : from(from), to(to), cost(cost), idx(idx) {}

  operator int() const { return to; }
};

template <typename T = int>
struct Graph {
  vector<vector<Edge<T> > > g;
  int es;

  Graph() = default;

  explicit Graph(int n) : g(n), es(0) {}

  size_t size() const { return g.size(); }

  void add_directed_edge(int from, int to, T cost = 1) {
    g[from].emplace_back(from, to, cost, es++);
  }

  void add_edge(int from, int to, T cost = 1) {
    g[from].emplace_back(from, to, cost, es);
    g[to].emplace_back(to, from, cost, es++);
  }

  void read(int M, int padding = -1, bool weighted = false,
            bool directed = false) {
    for (int i = 0; i < M; i++) {
      int a, b;
      cin >> a >> b;
      a += padding;
      b += padding;
      T c = T(1);
      if (weighted) cin >> c;
      if (directed)
        add_directed_edge(a, b, c);
      else
        add_edge(a, b, c);
    }
  }

  inline vector<Edge<T> > &operator[](const int &k) { return g[k]; }

  inline const vector<Edge<T> > &operator[](const int &k) const { return g[k]; }
};

template <typename T = int>
using Edges = vector<Edge<T> >;
#line 4 "graph/tree/centroid-decomposition.hpp"

/**
 * @brief Centroid-Decomosition(重心分解)
 */
template <typename T>
struct CentroidDecomposition : Graph<T> {
 public:
  using Graph<T>::Graph;
  using Graph<T>::g;
  Graph<int> tree;

  int build(int t = 0) {
    sub.assign(g.size(), 0);
    v.assign(g.size(), 0);
    tree = Graph<int>(g.size());
    return build_dfs(0);
  }

  explicit CentroidDecomposition(const Graph<T> &g) : Graph<T>(g) {}

 private:
  vector<int> sub;
  vector<int> v;

  inline int build_dfs(int idx, int par) {
    sub[idx] = 1;
    for (auto &to : g[idx]) {
      if (to == par || v[to]) continue;
      sub[idx] += build_dfs(to, idx);
    }
    return sub[idx];
  }

  inline int search_centroid(int idx, int par, const int mid) {
    for (auto &to : g[idx]) {
      if (to == par || v[to]) continue;
      if (sub[to] > mid) return search_centroid(to, idx, mid);
    }
    return idx;
  }

  inline int build_dfs(int idx) {
    int centroid = search_centroid(idx, -1, build_dfs(idx, -1) / 2);
    v[centroid] = true;
    for (auto &to : g[centroid]) {
      if (!v[to]) tree.add_directed_edge(centroid, build_dfs(to));
    }
    v[centroid] = false;
    return centroid;
  }
};
#line 6 "test/verify/yukicoder-1002.test.cpp"

int main() {
  int N, K;
  cin >> N >> K;
  CentroidDecomposition< int > g(N);
  g.read(N - 1, -1, true);
  int root = g.build();
  int64 ret = 0;
  vector< int > used(N);

  map< pair< int, int >, int > mp;
  int all;
  map< int, int > mp2;
  map< int, int > mp3;


  auto get_vec = MFP([&](auto get_vec, int idx, int par, int a, int b) -> void {
    if(b == -1) {
      ret += all - mp2[a];
      ret += mp3[a];
    } else {
      ret++;
      ret += mp2[a] + mp2[b];
      ret += mp[minmax(a, b)];
    }
    for(auto &e : g.g[idx]) {
      if(e.to == par) continue;
      if(used[e.to]) continue;
      if(a == e.cost) {
        get_vec(e.to, idx, e.cost, b);
      } else if(b == -1 || b == e.cost) {
        get_vec(e.to, idx, a, e.cost);
      }
    }
  });

  auto add_vec = MFP([&](auto add_vec, int idx, int par, int a, int b) -> void {
    if(b == -1) {
      mp2[a]++;
      all++;
    } else {
      mp[minmax(a, b)]++;
      mp3[a]++;
      mp3[b]++;
    }
    for(auto &e : g.g[idx]) {
      if(e.to == par) continue;
      if(used[e.to]) continue;
      if(a == e.cost) {
        add_vec(e.to, idx, e.cost, b);
      } else if(b == -1 || b == e.cost) {
        add_vec(e.to, idx, a, e.cost);
      }
    }
  });

  auto rec = MFP([&](auto rec, int idx) -> void {
    used[idx] = true;
    mp.clear();
    mp2.clear();
    mp3.clear();
    all = 0;
    for(auto &e : g.g[idx]) {
      if(used[e.to]) continue;
      get_vec(e.to, idx, e.cost, -1);
      add_vec(e.to, idx, e.cost, -1);
    }
    for(auto &to : g.tree.g[idx]) {
      rec(to);
    }
    used[idx] = false;
  });
  rec(root);
  cout << ret << endl;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	00_sample_00.in	AC	6 ms	4 MB
g++	00_sample_01.in	AC	6 ms	4 MB
g++	10_case_00.in	AC	6 ms	4 MB
g++	11_case_00.in	AC	172 ms	24 MB
g++	12_case_00.in	AC	237 ms	30 MB
g++	12_case_01.in	AC	236 ms	30 MB
g++	20_case_00.in	AC	7 ms	4 MB
g++	21_case_00.in	AC	141 ms	20 MB
g++	22_case_00.in	AC	261 ms	30 MB
g++	22_case_01.in	AC	247 ms	30 MB
g++	30_case_00.in	AC	7 ms	4 MB
g++	31_case_00.in	AC	265 ms	25 MB
g++	32_case_00.in	AC	380 ms	31 MB
g++	32_case_01.in	AC	372 ms	31 MB
g++	40_case_00.in	AC	6 ms	4 MB
g++	41_case_00.in	AC	170 ms	19 MB
g++	42_case_00.in	AC	348 ms	31 MB
g++	42_case_01.in	AC	381 ms	31 MB
g++	50_case_00.in	AC	7 ms	4 MB
g++	51_case_00.in	AC	314 ms	30 MB
g++	52_case_00.in	AC	362 ms	37 MB
g++	52_case_01.in	AC	376 ms	36 MB
g++	60_case_00.in	AC	7 ms	4 MB
g++	61_case_00.in	AC	311 ms	30 MB
g++	62_case_00.in	AC	637 ms	46 MB
g++	62_case_01.in	AC	604 ms	46 MB
g++	70_case_00.in	AC	6 ms	4 MB
g++	71_case_00.in	AC	50 ms	24 MB
g++	72_case_00.in	AC	80 ms	31 MB
g++	72_case_01.in	AC	73 ms	31 MB
g++	80_case_00.in	AC	7 ms	4 MB
g++	81_case_00.in	AC	69 ms	31 MB
g++	82_case_00.in	AC	76 ms	30 MB
g++	82_case_01.in	AC	70 ms	30 MB
g++	90_hand_01.in	AC	574 ms	55 MB
clang++	00_sample_00.in	AC	6 ms	4 MB
clang++	00_sample_01.in	AC	6 ms	4 MB
clang++	10_case_00.in	AC	6 ms	4 MB
clang++	11_case_00.in	AC	165 ms	24 MB
clang++	12_case_00.in	AC	252 ms	30 MB
clang++	12_case_01.in	AC	241 ms	30 MB
clang++	20_case_00.in	AC	7 ms	4 MB
clang++	21_case_00.in	AC	125 ms	20 MB
clang++	22_case_00.in	AC	238 ms	30 MB
clang++	22_case_01.in	AC	246 ms	30 MB
clang++	30_case_00.in	AC	7 ms	4 MB
clang++	31_case_00.in	AC	244 ms	25 MB
clang++	32_case_00.in	AC	390 ms	31 MB
clang++	32_case_01.in	AC	392 ms	31 MB
clang++	40_case_00.in	AC	6 ms	4 MB
clang++	41_case_00.in	AC	169 ms	19 MB
clang++	42_case_00.in	AC	405 ms	31 MB
clang++	42_case_01.in	AC	356 ms	31 MB
clang++	50_case_00.in	AC	7 ms	4 MB
clang++	51_case_00.in	AC	300 ms	31 MB
clang++	52_case_00.in	AC	370 ms	39 MB
clang++	52_case_01.in	AC	381 ms	39 MB
clang++	60_case_00.in	AC	7 ms	4 MB
clang++	61_case_00.in	AC	308 ms	29 MB
clang++	62_case_00.in	AC	628 ms	44 MB
clang++	62_case_01.in	AC	614 ms	44 MB
clang++	70_case_00.in	AC	6 ms	4 MB
clang++	71_case_00.in	AC	55 ms	24 MB
clang++	72_case_00.in	AC	89 ms	31 MB
clang++	72_case_01.in	AC	79 ms	30 MB
clang++	80_case_00.in	AC	7 ms	4 MB
clang++	81_case_00.in	AC	69 ms	30 MB
clang++	82_case_00.in	AC	76 ms	31 MB
clang++	82_case_01.in	AC	68 ms	31 MB
clang++	90_hand_01.in	AC	660 ms	54 MB

Back to top page