結果
| 問題 |
No.1002 Twotone
|
| コンテスト | |
| ユーザー |
 nono00
|
| 提出日時 | 2024-02-27 03:56:37 |
| 言語 | C++23 (gcc 13.3.0 + boost 1.87.0) |
| 結果 |
AC
|
| 実行時間 | 848 ms / 5,000 ms |
| コード長 | 7,324 bytes |
| コンパイル時間 | 1,927 ms |
| コンパイル使用メモリ | 146,892 KB |
| 実行使用メモリ | 33,780 KB |
| 最終ジャッジ日時 | 2024-09-29 11:53:54 |
| 合計ジャッジ時間 | 13,970 ms |
|
ジャッジサーバーID （参考情報） |
judge5 / judge2 |
（要ログイン）
| ファイルパターン | 結果 |
|---|---|
| sample | AC * 2 |
| other | AC * 33 |
ソースコード
#include <iostream>
#include <map>
#include <vector>
#include <iterator>
#include <ranges>
#include <vector>
namespace nono {
template <class T>
struct EdgeBase {
int from;
int to;
T weight;
EdgeBase() {}
EdgeBase(int from, int to, T weight = 1): from(from), to(to), weight(weight) {}
};
using Edge = EdgeBase<int>;
template <class T>
using WeightedEdge = EdgeBase<T>;
template <class T>
class Graph {
struct Edge_ {
int to;
T weight;
int id;
};
using iterator = std::vector<Edge_>::iterator;
using const_iterator = std::vector<Edge_>::const_iterator;
using subrange = std::ranges::subrange<iterator, iterator>;
using const_subrange = std::ranges::subrange<const_iterator, const_iterator>;
public:
template <class U>
friend Graph<U> to_undirected_graph(int n, const std::vector<EdgeBase<U>>& edges);
template <class U>
friend Graph<U> to_directed_graph(int n, const std::vector<EdgeBase<U>>& edges);
subrange operator[](int i) {
return std::ranges::subrange(edges_.begin() + indptr_[i], edges_.begin() + indptr_[i + 1]);
}
const_subrange operator[](int i) const {
return std::ranges::subrange(edges_.begin() + indptr_[i], edges_.begin() + indptr_[i + 1]);
}
int size() const {
return n_;
}
bool is_directed() const {
return directed_;
}
bool is_undirected() const {
return !is_directed();
}
private:
Graph(int n, const std::vector<EdgeBase<T>>& edges, bool directed)
: n_(n),
indptr_(n_ + 1),
edges_(directed ? edges.size() : 2 * edges.size()),
directed_(directed) {
for (const auto& e: edges) {
indptr_[e.from + 1]++;
if (!directed_) indptr_[e.to + 1]++;
}
for (int i = 0; i < n_; i++) {
indptr_[i + 1] += indptr_[i];
}
auto index = indptr_;
for (int i = 0; i < std::ssize(edges); i++) {
const auto& e = edges[i];
edges_[index[e.from]++] = Edge_(e.to, e.weight, i);
if (!directed_) edges_[index[e.to]++] = Edge_(e.from, e.weight, i);
}
}
int n_;
std::vector<int> indptr_;
std::vector<Edge_> edges_;
bool directed_;
};
template <class T>
Graph<T> to_undirected_graph(int n, const std::vector<EdgeBase<T>>& edges) {
return Graph<T>(n, edges, false);
}
template <class T>
Graph<T> to_directed_graph(int n, const std::vector<EdgeBase<T>>& edges) {
return Graph<T>(n, edges, true);
}
} // namespace nono
#include <vector>
#include <algorithm>
#include <vector>
namespace nono {
template <class T>
bool is_tree(const Graph<T>& graph) {
if (graph.is_directed()) return false;
constexpr int NONE = -1;
int n = graph.size();
std::vector<int> used(n);
auto dfs = [&](auto self, int u, int eid) -> bool {
used[u] = 1;
for (const auto& e: graph[u]) {
if (e.id == eid) continue;
if (used[e.to]) return false;
if (!self(self, e.to, e.id)) return false;
}
return true;
};
if (!dfs(dfs, 0, NONE)) return false;
return std::ranges::all_of(used, [](int f) { return f == 1; });
}
} // namespace nono
namespace nono {
template <class T>
std::vector<int> centroids(const Graph<T>& graph) {
int n = graph.size();
std::vector<int> subtree(n);
std::vector<bool> removed(n);
std::vector<int> result;
result.reserve(n);
auto calc_subtree = [&](auto&& self, int u, int p) -> int {
subtree[u] = 1;
for (const auto& e: graph[u]) {
if (e.to == p || removed[e.to]) continue;
subtree[u] += self(self, e.to, u);
}
return subtree[u];
};
std::vector<int> cur;
cur.push_back(0);
while (!cur.empty()) {
std::vector<int> next;
for (auto root: cur) {
calc_subtree(calc_subtree, root, root);
int centroid = root;
bool flag = true;
while (flag) {
flag = false;
for (const auto& e: graph[centroid]) {
if (removed[e.to] || subtree[e.to] > subtree[centroid]) continue;
if (subtree[root] <= 2 * subtree[e.to]) {
centroid = e.to;
flag = true;
}
}
}
result.push_back(centroid);
removed[centroid] = true;
for (const auto& e: graph[centroid]) {
if (removed[e.to]) continue;
next.push_back(e.to);
}
}
cur = std::move(next);
}
return result;
}
} // namespace nono
namespace nono {
void solve() {
int n, k;
std::cin >> n >> k;
std::vector<WeightedEdge<int>> edges;
edges.reserve(n - 1);
for (int i = 0; i + 1 < n; i++) {
int u, v, c;
std::cin >> u >> v >> c;
u--;
v--;
edges.emplace_back(u, v, c);
}
const auto graph = to_undirected_graph(n, edges);
std::vector<bool> removed(n);
long long ans = 0;
for (auto centroid: centroids(graph)) {
std::map<std::pair<int, int>, int> a;
std::map<int, int> b;
std::map<int, int> c;
int d = 0;
std::vector<std::pair<int, int>> path;
auto dfs = [&](auto&& self, int u, int p, std::pair<int, int> color) -> void {
path.push_back(color);
for (const auto& e: graph[u]) {
if (e.to == p || removed[e.to]) continue;
auto new_color = color;
if (new_color.first == -1) {
new_color.first = e.weight;
self(self, e.to, u, new_color);
} else if (new_color.first != e.weight && new_color.second == -1) {
new_color.second = e.weight;
self(self, e.to, u, new_color);
} else if (new_color.first == e.weight || new_color.second == e.weight) {
self(self, e.to, u, new_color);
}
}
};
for (const auto& e: graph[centroid]) {
if (removed[e.to]) continue;
path.clear();
dfs(dfs, e.to, centroid, {e.weight, -1});
for (auto [c1, c2]: path) {
if (c1 == -1) continue;
if (c2 == -1) {
ans += b[c1];
ans += d - c[c1];
} else {
if (c1 > c2) std::swap(c1, c2);
ans += a[std::pair(c1, c2)];
ans += c[c1];
ans += c[c2];
ans++;
}
}
for (auto [c1, c2]: path) {
if (c1 == -1) continue;
if (c2 == -1) {
c[c1]++;
d++;
} else {
if (c1 > c2) std::swap(c1, c2);
a[std::pair(c1, c2)]++;
b[c1]++;
b[c2]++;
}
}
}
removed[centroid] = true;
}
std::cout << ans << '\n';
}
} // namespace nono
int main() {
std::cin.tie(0)->sync_with_stdio(0);
nono::solve();
}