#define make_pair mp #define emplace_back pb #include <bits/stdc++.h> using namespace std; mt19937 mt_rand(time(0)); const int N = 25005; const long long MOD = 1e9 + 7; int n, k; long long pots[N]; vector< pair<int, int> > v[N]; vector< vector<int> > answers; bool inv[N]; int maksp = 0; vector<int> tmp; void dfs(int x, int y) { if(x < y) answers.pb(tmp); inv[x] = 1; for(auto p : v[x]) { int pp = p.second; int z = p.first; if(!inv[z]) { tmp[pp - 1]++; dfs(z, y); tmp[pp - 1]--; } } } int cnt[N]; vector<int> find(int k, int pos, vector< vector<int> > &res) { /*cout<<"KOLEJNE!"<<" "<<k<<" "<<pos<<endl; for(auto vec : res) { for(auto x : vec) cout<<x<<" "; cout<<endl; }*/ for(int i=0;i<=n;i++) cnt[i] = 0; for(auto &vec : res) cnt[vec[pos]]++; int wsk = 0; int sum = 0; for(int i=0;i<n;i++) { sum += cnt[i]; if(sum >= k) { wsk = i; sum -= cnt[i]; break; } } if(pos == 0) for(auto &vec : res) if(vec[pos] == wsk) return vec; vector< vector<int> > res2; for(auto &vec : res) if(vec[pos] == wsk) { res2.pb(vec); } return find(k - sum, pos - 1, res2); } int main() { scanf("%d%d", &n, &k); pots[0] = 1; for(int i=1;i<=n;i++) pots[i] = pots[i-1] * (long long)n % MOD; for(int i=1;i<n;i++) { int a, b, p; scanf("%d%d%d", &a, &b, &p); v[a].pb(mp(b, p)); v[b].pb(mp(a, p)); maksp = max(maksp, p); } tmp.resize(maksp); for(int i=2;i<=n;i++) { dfs(i, i); for(int j=1;j<=n;j++) inv[j] = 0; } auto res = find(k, maksp - 1, answers); long long ans = 0; for(int i=0;i<maksp;i++) { long long c = res[i]; ans = (ans + c * pots[i + 1]) % MOD; } ans %= MOD; printf("%lld", ans); return 0; }
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 77 78 79 80 81 82 83 84 85 86 87 88 | #define make_pair mp #define emplace_back pb #include <bits/stdc++.h> using namespace std; mt19937 mt_rand(time(0)); const int N = 25005; const long long MOD = 1e9 + 7; int n, k; long long pots[N]; vector< pair<int, int> > v[N]; vector< vector<int> > answers; bool inv[N]; int maksp = 0; vector<int> tmp; void dfs(int x, int y) { if(x < y) answers.pb(tmp); inv[x] = 1; for(auto p : v[x]) { int pp = p.second; int z = p.first; if(!inv[z]) { tmp[pp - 1]++; dfs(z, y); tmp[pp - 1]--; } } } int cnt[N]; vector<int> find(int k, int pos, vector< vector<int> > &res) { /*cout<<"KOLEJNE!"<<" "<<k<<" "<<pos<<endl; for(auto vec : res) { for(auto x : vec) cout<<x<<" "; cout<<endl; }*/ for(int i=0;i<=n;i++) cnt[i] = 0; for(auto &vec : res) cnt[vec[pos]]++; int wsk = 0; int sum = 0; for(int i=0;i<n;i++) { sum += cnt[i]; if(sum >= k) { wsk = i; sum -= cnt[i]; break; } } if(pos == 0) for(auto &vec : res) if(vec[pos] == wsk) return vec; vector< vector<int> > res2; for(auto &vec : res) if(vec[pos] == wsk) { res2.pb(vec); } return find(k - sum, pos - 1, res2); } int main() { scanf("%d%d", &n, &k); pots[0] = 1; for(int i=1;i<=n;i++) pots[i] = pots[i-1] * (long long)n % MOD; for(int i=1;i<n;i++) { int a, b, p; scanf("%d%d%d", &a, &b, &p); v[a].pb(mp(b, p)); v[b].pb(mp(a, p)); maksp = max(maksp, p); } tmp.resize(maksp); for(int i=2;i<=n;i++) { dfs(i, i); for(int j=1;j<=n;j++) inv[j] = 0; } auto res = find(k, maksp - 1, answers); long long ans = 0; for(int i=0;i<maksp;i++) { long long c = res[i]; ans = (ans + c * pots[i + 1]) % MOD; } ans %= MOD; printf("%lld", ans); return 0; } |