#include <bits/stdc++.h>
using namespace std;
using ll = long long int;
const int MAXN = 25010;
vector<pair<int, int>> graf[MAXN];
bitset<MAXN> pairs[MAXN];
int am[MAXN];
const int MOD = 1000 * 1000 * 1000 + 7;
ll pot[MAXN];
int req = 0;
int root = 0;
int actp = 0;
void DFS(int v, int prevv, int cnt) {
if (!pairs[root].test(v)) am[cnt]++;
for (pair<int, int> p : graf[v])
if (p.first != prevv)
DFS(p.first, v, cnt + (p.second == actp));
}
void DFSE(int v, int prevv, int cnt) {
if (cnt != req) pairs[root].set(v, true);
for (pair<int, int> p : graf[v])
if (p.first != prevv)
DFSE(p.first, v, cnt + (p.second == actp));
}
int main() {
int n;
ll k;
scanf("%d%lld", &n, &k);
k = 2 * k + n;
int maxp = 0;
for (int i = 1; i < n; i++) {
int a, b, p;
scanf("%d%d%d", &a, &b, &p);
graf[a].push_back(make_pair(b, p));
graf[b].push_back(make_pair(a, p));
maxp = max(maxp, p);
}
pot[1] = n;
for (int i = 2; i <= maxp; i++) pot[i] = (pot[i - 1] * n) % MOD;
ll out = 0;
for (actp = maxp; actp >= 1; actp--) {
for (int i = 0; i < n; i++) am[i] = 0;
for (int i = 1; i <= n; i++) {
root = i;
DFS(i, i, 0);
}
req = 0;
while (am[req] < k) {
k -= am[req];
req++;
out += pot[actp];
if (out >= MOD) out -= MOD;
}
for (int i = 1; i <= n; i++) {
root = i;
DFSE(i, i, 0);
}
}
printf("%lld", out);
}
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | #include <bits/stdc++.h> using namespace std; using ll = long long int; const int MAXN = 25010; vector<pair<int, int>> graf[MAXN]; bitset<MAXN> pairs[MAXN]; int am[MAXN]; const int MOD = 1000 * 1000 * 1000 + 7; ll pot[MAXN]; int req = 0; int root = 0; int actp = 0; void DFS(int v, int prevv, int cnt) { if (!pairs[root].test(v)) am[cnt]++; for (pair<int, int> p : graf[v]) if (p.first != prevv) DFS(p.first, v, cnt + (p.second == actp)); } void DFSE(int v, int prevv, int cnt) { if (cnt != req) pairs[root].set(v, true); for (pair<int, int> p : graf[v]) if (p.first != prevv) DFSE(p.first, v, cnt + (p.second == actp)); } int main() { int n; ll k; scanf("%d%lld", &n, &k); k = 2 * k + n; int maxp = 0; for (int i = 1; i < n; i++) { int a, b, p; scanf("%d%d%d", &a, &b, &p); graf[a].push_back(make_pair(b, p)); graf[b].push_back(make_pair(a, p)); maxp = max(maxp, p); } pot[1] = n; for (int i = 2; i <= maxp; i++) pot[i] = (pot[i - 1] * n) % MOD; ll out = 0; for (actp = maxp; actp >= 1; actp--) { for (int i = 0; i < n; i++) am[i] = 0; for (int i = 1; i <= n; i++) { root = i; DFS(i, i, 0); } req = 0; while (am[req] < k) { k -= am[req]; req++; out += pot[actp]; if (out >= MOD) out -= MOD; } for (int i = 1; i <= n; i++) { root = i; DFSE(i, i, 0); } } printf("%lld", out); } |
English