#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); } |