#line 1 "test/yukicoder/1002.test.cpp"
#define PROBLEM "https://yukicoder.me/problems/no/1002"
#include <bits/stdc++.h>
#line 1 "library/tree/CentroidDecomposition.hpp"
template <typename TREE> class CentroidDecomposition {
TREE T;
std::vector<int> sz, pre, timing;
int find_centroid(int v) {
std::vector<int> S{v};
pre[v] = -1;
for (int i = 0; i < S.size(); i++) {
const int u = S[i];
sz[u] = 1;
for (int to : T[u]) {
if (to == pre[u] || ~timing[to])
continue;
pre[to] = u;
S.push_back(to);
}
}
int SZ = S.size();
std::ranges::reverse(S);
for (int u : S) {
if (SZ - sz[u] <= SZ / 2)
return u;
sz[pre[u]] += sz[u];
}
assert(false);
return -1;
};
public:
std::vector<int> order;
CentroidDecomposition(TREE T) : T(T), sz(T.n), pre(T.n), timing(T.n, -1) {
order.reserve(T.n);
std::queue<int> que;
que.push(0);
while (que.size()) {
int c = find_centroid(que.front());
que.pop();
timing[c] = order.size();
order.push_back(c);
for (int to : T[c])
if (timing[to] < 0)
que.push(to);
}
}
template <typename X, typename F, typename G, typename H>
void calc(int root, X initial_val, const F &next_val, const G &action,
const H &finish) {
std::queue<std::tuple<int, int, X>> que;
auto f = [&](int v_, int pre_, X val_, bool is_all) {
que.emplace(v_, pre_, val_);
while (que.size()) {
auto [v, pre, val] = que.front();
que.pop();
action(val, is_all);
for (const auto &e : T[v]) {
if (e.to == pre || timing[e.to] <= timing[root])
continue;
que.emplace(e.to, v, next_val(val, e));
}
}
finish(is_all);
};
for (const auto &e : T[root])
if (timing[e.to] > timing[root])
f(e.to, root, next_val(initial_val, e), false);
f(root, -1, initial_val, true);
}
template <typename X, typename F, typename G, typename H>
void all_calc(X initial_val, const F &next_val, const G &action,
const H &finish) {
for (int i = 0; i < T.n; i++)
calc(i, initial_val, next_val, action, finish);
}
};
#line 2 "library/graph/WeightedGraph.hpp"
template <typename T> struct WeightedEdge {
WeightedEdge() = default;
WeightedEdge(int from, int to, T weight)
: from(from), to(to), weight(weight) {}
int from, to;
T weight;
operator int() const { return to; }
};
template <typename T> struct WeightedGraph {
int n;
using weight_type = T;
using edge_type = WeightedEdge<T>;
std::vector<edge_type> edges;
protected:
std::vector<int> in_deg;
bool prepared;
class OutgoingEdges {
WeightedGraph *g;
int l, r;
public:
OutgoingEdges(WeightedGraph *g, int l, int r) : g(g), l(l), r(r) {}
edge_type *begin() { return &(g->edges[l]); }
edge_type *end() { return &(g->edges[r]); }
edge_type &operator[](int i) { return g->edges[l + i]; }
int size() const { return r - l; }
};
class ConstOutgoingEdges {
const WeightedGraph *g;
int l, r;
public:
ConstOutgoingEdges(const WeightedGraph *g, int l, int r)
: g(g), l(l), r(r) {}
const edge_type *begin() const { return &(g->edges[l]); }
const edge_type *end() const { return &(g->edges[r]); }
const edge_type &operator[](int i) const { return g->edges[l + i]; }
int size() const { return r - l; }
};
public:
OutgoingEdges operator[](int v) {
assert(prepared);
return {this, in_deg[v], in_deg[v + 1]};
}
const ConstOutgoingEdges operator[](int v) const {
assert(prepared);
return {this, in_deg[v], in_deg[v + 1]};
}
bool is_prepared() const { return prepared; }
WeightedGraph() : n(0), in_deg(1, 0), prepared(false) {}
WeightedGraph(int n) : n(n), in_deg(n + 1, 0), prepared(false) {}
WeightedGraph(int n, int m, bool directed = false, int indexed = 1)
: n(n), in_deg(n + 1, 0), prepared(false) {
scan(m, directed, indexed);
}
void resize(int n) { n = n; }
void add_arc(int from, int to, T weight) {
assert(!prepared);
assert(0 <= from and from < n and 0 <= to and to < n);
edges.emplace_back(from, to, weight);
in_deg[from + 1]++;
}
void add_edge(int u, int v, T weight) {
add_arc(u, v, weight);
add_arc(v, u, weight);
}
void add_arc(const edge_type &e) { add_arc(e.from, e.to, e.weight); }
void add_edge(const edge_type &e) { add_edge(e.from, e.to, e.weight); }
void scan(int m, bool directed = false, int indexed = 1) {
edges.reserve(directed ? m : 2 * m);
while (m--) {
int u, v;
std::cin >> u >> v;
u -= indexed;
v -= indexed;
T weight;
std::cin >> weight;
if (directed)
add_arc(u, v, weight);
else
add_edge(u, v, weight);
}
build();
}
void build() {
assert(!prepared);
prepared = true;
for (int v = 0; v < n; v++)
in_deg[v + 1] += in_deg[v];
std::vector<edge_type> new_edges(in_deg.back());
auto counter = in_deg;
for (auto &&e : edges)
new_edges[counter[e.from]++] = e;
edges = new_edges;
}
void graph_debug() const {
#ifndef __DEBUG
return;
#endif
assert(prepared);
for (int from = 0; from < n; from++) {
std::cerr << from << ";";
for (int i = in_deg[from]; i < in_deg[from + 1]; i++)
std::cerr << "(" << edges[i].to << "," << edges[i].weight
<< ")";
std::cerr << "\n";
}
}
};
#line 3 "library/tree/WeightedTree.hpp"
template <typename T> struct WeightedTree : WeightedGraph<T> {
using WeightedGraph<T>::WeightedGraph;
using edge_type = typename WeightedGraph<T>::edge_type;
using OutgoingEdges = typename WeightedGraph<T>::OutgoingEdges;
using WeightedGraph<T>::n;
using WeightedGraph<T>::in_deg;
int root = -1;
std::vector<int> DFS, BFS, depth;
void scan_root(int indexed = 1) {
for (int i = 1; i < n; i++) {
int p;
std::cin >> p;
T weight;
std::cin >> weight;
add_edge(p - indexed, i, weight);
}
build();
}
void scan(int indexed = 1) {
WeightedGraph<T>::scan(n - 1, false, indexed);
build();
}
edge_type &parent(int v) {
assert(~root and root != v);
return (*this)[v][0];
}
OutgoingEdges son(int v) {
assert(~root);
if (v == root)
return {this, in_deg[v], in_deg[v + 1]};
return {this, in_deg[v] + 1, in_deg[v + 1]};
}
private:
void dfs(int v, int pre = -1) {
for (auto &e : (*this)[v]) {
if (e.to == pre)
std::swap((*this)[v][0], e);
else {
depth[e.to] = depth[v] + 1;
dfs(e.to, v);
}
}
DFS.push_back(v);
}
public:
void build(int r = 0) {
if (!WeightedGraph<T>::is_prepared())
WeightedGraph<T>::build();
if (~root) {
assert(r == root);
return;
}
root = r;
depth = std::vector<int>(n, 0);
DFS.reserve(n);
BFS.reserve(n);
dfs(root);
std::queue<int> que;
que.push(root);
while (que.size()) {
int p = que.front();
que.pop();
BFS.push_back(p);
for (const auto &e : son(p))
que.push(e.to);
}
}
};
#line 6 "test/yukicoder/1002.test.cpp"
using ll = long long;
int main() {
std::ios::sync_with_stdio(false);
std::cin.tie(nullptr);
int n, k;
std::cin >> n >> k;
WeightedTree<int> T(n);
T.scan(1);
ll ans = 0;
CentroidDecomposition CD(T);
std::map<std::pair<int, int>, int> m2; // m2[c1,c2]:c1,c2 組の数
std::map<int, int> m1, // m1[c]:c 単色の数
m2_sum; // m2_sum[c]:c を含む二色の数
int m1_sum = 0; // 単色の合計
bool empty = false;
// k+1 : 色が三色以上
// -1 : 色が未定
// 色は後ろから埋まる 両方埋まった場合は minmax
auto F = [&](const std::pair<int, int> &cc, const auto &e) {
auto [c1, c2] = cc;
int c = e.weight;
if (c1 > k || c1 == c || c2 == c)
return cc;
if (~c1)
return std::make_pair(k + 1, k + 1);
if (c > c2)
std::swap(c, c2);
return std::make_pair(c, c2);
};
auto G = [&](std::pair<int, int> cc, bool add) {
auto [c1, c2] = cc;
if (c1 > k)
return;
ll sum = 0;
if (~c1) {
if (m2.count(cc))
sum += m2[cc];
if (m1.count(c1))
sum += m1[c1];
if (m1.count(c2))
sum += m1[c2];
if (empty)
sum++;
m2[cc]++;
m2_sum[c1]++;
m2_sum[c2]++;
} else if (~c2) {
sum += m1_sum;
if (m1.count(c2))
sum -= m1[c2];
if (m2_sum.count(c2))
sum += m2_sum[c2];
m1_sum++;
m1[c2]++;
} else
empty = true;
if (add)
ans += sum;
else
ans -= sum;
};
auto H = [&](bool add) {
m2.clear();
m1.clear();
m2_sum.clear();
m1_sum = empty = 0;
};
CD.all_calc(std::make_pair(-1, -1), F, G, H);
std::cout << ans << std::endl;
}
test/yukicoder/1002.test.cpp
build/pch/stdc++.hpp