#include <bits/stdc++.h> using namespace std; const int MX=25005,BIG=82000000,md=1000000007; uint16_t c[BIG*4]; int st[BIG],n,m,k,i,j,e,v,x,y,z[MX],q[MX],fi,fr,trnew,trmax,tr[MX],cnt[MX],pdm,idm,id,r; vector<int> g[MX],f[MX]; long long pw[MX]; int main() { scanf("%d%d",&n,&k); for (i=1; i<n; i++) { scanf("%d%d%d",&x,&y,&z[i-1]); g[x].push_back(y); f[x].push_back(z[i-1]); g[y].push_back(x); f[y].push_back(z[i-1]); } sort(z,z+n-1); m=unique(z,z+n-1)-z; for (pw[0]=i=1; i<=n; i++) { pw[i]=(pw[i-1]*n)%md; tr[i]=((i-2)*(i-1))/2; for (j=0; j<f[i].size(); j++) { f[i][j]=lower_bound(z,z+m-1,f[i][j])-z; x=g[i][j]; if (x<i) { id=tr[i]+x; st[id]=i; c[id*m+f[i][j]]=1; q[fr++]=id; } } } for (fi=0; fi<fr; fi++) { i=q[fi]; x=st[i]; y=i-tr[x]; pdm=i*m; for (j=0; j<g[y].size(); j++) if ((v=g[y][j])<x) { id=tr[x]+v; if (st[id]) continue; st[id]=x; for (idm=id*m, e=0; e<m; e++) c[idm+e]=c[pdm+e]; ++c[idm+f[y][j]]; q[fr++]=id; } for (j=0; j<g[x].size(); j++) if ((v=g[x][j])<y) { id=tr[y]+v; if (st[id]) continue; st[id]=y; for (idm=id*m, e=0; e<m; e++) c[idm+e]=c[pdm+e]; ++c[idm+f[x][j]]; q[fr++]=id; } } trmax=(n*(n-1))/2; for (i=0; i<trmax; i++) st[i]=i+1; for (j=m-1; j>=0; --j) { if (j<m-1) memset(cnt,0,sizeof(cnt)); for (i=0; i<trmax; ++i) ++cnt[c[st[i]*m+j]]; for (e=0; e<=n; ++e) { if (cnt[e]>=k) break; k-=cnt[e]; } r=(r+e*pw[z[j]])%md; for (trnew=i=0; i<trmax; ++i) if (c[st[i]*m+j]==e) st[trnew++]=st[i]; trmax=trnew; } printf("%d\n",r); 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 | #include <bits/stdc++.h> using namespace std; const int MX=25005,BIG=82000000,md=1000000007; uint16_t c[BIG*4]; int st[BIG],n,m,k,i,j,e,v,x,y,z[MX],q[MX],fi,fr,trnew,trmax,tr[MX],cnt[MX],pdm,idm,id,r; vector<int> g[MX],f[MX]; long long pw[MX]; int main() { scanf("%d%d",&n,&k); for (i=1; i<n; i++) { scanf("%d%d%d",&x,&y,&z[i-1]); g[x].push_back(y); f[x].push_back(z[i-1]); g[y].push_back(x); f[y].push_back(z[i-1]); } sort(z,z+n-1); m=unique(z,z+n-1)-z; for (pw[0]=i=1; i<=n; i++) { pw[i]=(pw[i-1]*n)%md; tr[i]=((i-2)*(i-1))/2; for (j=0; j<f[i].size(); j++) { f[i][j]=lower_bound(z,z+m-1,f[i][j])-z; x=g[i][j]; if (x<i) { id=tr[i]+x; st[id]=i; c[id*m+f[i][j]]=1; q[fr++]=id; } } } for (fi=0; fi<fr; fi++) { i=q[fi]; x=st[i]; y=i-tr[x]; pdm=i*m; for (j=0; j<g[y].size(); j++) if ((v=g[y][j])<x) { id=tr[x]+v; if (st[id]) continue; st[id]=x; for (idm=id*m, e=0; e<m; e++) c[idm+e]=c[pdm+e]; ++c[idm+f[y][j]]; q[fr++]=id; } for (j=0; j<g[x].size(); j++) if ((v=g[x][j])<y) { id=tr[y]+v; if (st[id]) continue; st[id]=y; for (idm=id*m, e=0; e<m; e++) c[idm+e]=c[pdm+e]; ++c[idm+f[x][j]]; q[fr++]=id; } } trmax=(n*(n-1))/2; for (i=0; i<trmax; i++) st[i]=i+1; for (j=m-1; j>=0; --j) { if (j<m-1) memset(cnt,0,sizeof(cnt)); for (i=0; i<trmax; ++i) ++cnt[c[st[i]*m+j]]; for (e=0; e<=n; ++e) { if (cnt[e]>=k) break; k-=cnt[e]; } r=(r+e*pw[z[j]])%md; for (trnew=i=0; i<trmax; ++i) if (c[st[i]*m+j]==e) st[trnew++]=st[i]; trmax=trnew; } printf("%d\n",r); return 0; } |