結果
| 問題 |
No.1002 Twotone
|
| コンテスト | |
| ユーザー |
 yosupot
|
| 提出日時 | 2020-02-28 22:37:47 |
| 言語 | C++17 (gcc 13.3.0 + boost 1.87.0) |
| 結果 |
AC
|
| 実行時間 | 971 ms / 5,000 ms |
| コード長 | 7,403 bytes |
| コンパイル時間 | 2,054 ms |
| コンパイル使用メモリ | 147,532 KB |
| 最終ジャッジ日時 | 2025-01-09 02:58:51 |
|
ジャッジサーバーID （参考情報） |
judge3 / judge2 |
（要ログイン）
| ファイルパターン | 結果 |
|---|---|
| sample | AC * 2 |
| other | AC * 33 |
ソースコード
//#pragma GCC optimize("Ofast")
//#pragma GCC target("avx")
//#undef LOCAL
#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <complex>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
using namespace std;
using uint = unsigned int;
using ll = long long;
using ull = unsigned long long;
constexpr ll TEN(int n) { return (n == 0) ? 1 : 10 * TEN(n - 1); }
template <class T> using V = vector<T>;
template <class T> using VV = V<V<T>>;
struct Scanner {
FILE* fp = nullptr;
char line[(1 << 15) + 1];
size_t st = 0, ed = 0;
void reread() {
memmove(line, line + st, ed - st);
ed -= st;
st = 0;
ed += fread(line + ed, 1, (1 << 15) - ed, fp);
line[ed] = '\0';
}
bool succ() {
while (true) {
if (st == ed) {
reread();
if (st == ed) return false;
}
while (st != ed && isspace(line[st])) st++;
if (st != ed) break;
}
if (ed - st <= 50) reread();
return true;
}
template <class T, enable_if_t<is_same<T, string>::value, int> = 0>
bool read_single(T& ref) {
if (!succ()) return false;
while (true) {
size_t sz = 0;
while (st + sz < ed && !isspace(line[st + sz])) sz++;
ref.append(line + st, sz);
st += sz;
if (!sz || st != ed) break;
reread();
}
return true;
}
template <class T, enable_if_t<is_integral<T>::value, int> = 0>
bool read_single(T& ref) {
if (!succ()) return false;
bool neg = false;
if (line[st] == '-') {
neg = true;
st++;
}
ref = T(0);
while (isdigit(line[st])) {
ref = 10 * ref + (line[st++] - '0');
}
if (neg) ref = -ref;
return true;
}
template <class T> bool read_single(V<T>& ref) {
for (auto& d : ref) {
if (!read_single(d)) return false;
}
return true;
}
void read() {}
template <class H, class... T> void read(H& h, T&... t) {
bool f = read_single(h);
assert(f);
read(t...);
}
Scanner(FILE* _fp) : fp(_fp) {}
};
struct Printer {
public:
template <bool F = false> void write() {}
template <bool F = false, class H, class... T>
void write(const H& h, const T&... t) {
if (F) write_single(' ');
write_single(h);
write<true>(t...);
}
template <class... T> void writeln(const T&... t) {
write(t...);
write_single('\n');
}
Printer(FILE* _fp) : fp(_fp) {}
~Printer() { flush(); }
private:
static constexpr size_t SIZE = 1 << 15;
FILE* fp;
char line[SIZE], small[50];
size_t pos = 0;
void flush() {
fwrite(line, 1, pos, fp);
pos = 0;
}
void write_single(const char& val) {
if (pos == SIZE) flush();
line[pos++] = val;
}
template <class T, enable_if_t<is_integral<T>::value, int> = 0>
void write_single(T val) {
if (pos > (1 << 15) - 50) flush();
if (val == 0) {
write_single('0');
return;
}
if (val < 0) {
write_single('-');
val = -val; // todo min
}
size_t len = 0;
while (val) {
small[len++] = char('0' + (val % 10));
val /= 10;
}
for (size_t i = 0; i < len; i++) {
line[pos + i] = small[len - 1 - i];
}
pos += len;
}
void write_single(const string& s) {
for (char c : s) write_single(c);
}
void write_single(const char* s) {
size_t len = strlen(s);
for (size_t i = 0; i < len; i++) write_single(s[i]);
}
template <class T> void write_single(const V<T>& val) {
auto n = val.size();
for (size_t i = 0; i < n; i++) {
if (i) write_single(' ');
write_single(val[i]);
}
}
};
struct Centroid {
int r;
VV<int> tr;
V<int> par;
};
template <class E> struct CentroidExec : Centroid {
int n;
const VV<E>& g;
V<bool> used;
using P = pair<int, int>;
V<P> info; //(child max, child)
int dfs(int p, int b) {
int sz = 1;
info[p] = P(0, -1);
for (E e : g[p]) {
int d = e.to;
if (d == b || used[d]) continue;
int csz = dfs(d, p);
sz += csz;
info[p] = max(info[p], P(csz, d));
}
return sz;
}
int init(int p, int b) {
int sz = dfs(p, -1);
while (info[p].first > sz / 2) p = info[p].second;
par[p] = b;
used[p] = true;
for (E e : g[p]) {
int d = e.to;
if (used[d]) continue;
tr[p].push_back(init(d, p));
}
return p;
}
CentroidExec(const VV<E>& _g) : n(int(_g.size())), g(_g), used(n), info(n) {
tr = VV<int>(n);
par = V<int>(n);
r = init(0, -1);
}
};
template <class E> Centroid get_centroid(const VV<E>& g) {
return CentroidExec<E>(g);
}
Scanner sc = Scanner(stdin);
Printer pr = Printer(stdout);
int main() {
int n, k;
sc.read(n, k);
struct E {
int to, col;
};
VV<E> g(n);
for (int i = 0; i < n - 1; i++) {
int a, b, c;
sc.read(a, b, c); a--; b--;
g[a].push_back({b, c});
g[b].push_back({a, c});
}
auto cent = get_centroid(g);
V<bool> cvis(n);
using P = pair<int, int>;
ll ans = 0;
auto dfs = [&](auto self, int p) -> void {
int mp1_sum = 0;
map<int, int> mp1, mp2_sum;
map<P, int> mp2;
V<P> buf;
auto inc = [&](int c0, int c1) {
if (c0 < c1) swap(c0, c1);
if (c1 != -1) {
ans++;
ans += mp2[{c0, c1}] + mp1[c0] + mp1[c1];
} else {
ans += mp1_sum - mp1[c0] + mp2_sum[c0];
}
buf.push_back({c0, c1});
};
auto ref = [&]() {
for (auto p: buf) {
if (p.second != -1) {
mp2[p]++;
mp2_sum[p.first]++;
mp2_sum[p.second]++;
} else {
mp1[p.first]++;
mp1_sum++;
}
}
buf.clear();
};
auto dfs2 = [&](auto self, int p, int b, int c0, int c1) -> void {
inc(c0, c1);
for (auto e: g[p]) {
int d = e.to;
int c = e.col;
if (d == b || cvis[d]) continue;
if (c == c0 || c == c1) {
self(self, d, p, c0, c1);
} else if (c1 == -1) {
self(self, d, p, c0, c);
}
}
};
for (auto e: g[p]) {
int d = e.to;
if (cvis[d]) continue;
dfs2(dfs2, d, p, e.col, -1);
ref();
}
cvis[p] = true;
for (auto d: cent.tr[p]) {
self(self, d);
}
};
dfs(dfs, cent.r);
pr.writeln(ans);
return 0;
}