Advertisement
RaFiN_

how many pairs in a tree that ....

Apr 27th, 2019
135
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 2.26 KB | None | 0 0
  1. /* how many pairs in a tree that constitute a number that is divisible by M */
  2.  
  3. /* Cf - 716E */
  4.  
  5. int Subtree[MAX],del[MAX],n,totalnode;
  6.  
  7. ll inv[MAX],power[MAX];ll M;vector<pll>v[MAX];ll up[MAX],down[MAX],idx,deep[MAX];vi lst;ll ans;ll comp[MAX];
  8.  
  9. void precal(){
  10.     for(ll b = 1;b<=10;b++){
  11.         if(b*M %10 !=1)continue;
  12.         inv[1] =( 1 - b*M)/10 + M;
  13.         inv[1] %=  M;
  14.     }
  15.     for(int i = 2;i<MAX;i++)inv[i] = (inv[i-1]*inv[1]+M)%M;
  16.     power[0] = 1;
  17.     for(int i = 1;i<MAX;i++)power[i] = (power[i-1]*10)%M;
  18. }
  19.  
  20. void dfs2(int s,int p)
  21. {
  22.     Subtree[s] = 1;
  23.     totalnode++;
  24.     for(auto it : v[s])
  25.     {
  26.         if(it.ff==p||del[it.ff])continue;
  27.         dfs2(it.ff,s);
  28.         Subtree[s]+=Subtree[it.ff];
  29.     }
  30. }
  31.  
  32. int Getcenter(int s,int p)
  33. {
  34.     for(auto it : v[s])
  35.     {
  36.         if(it.ff==p||del[it.ff])continue;
  37.         if(Subtree[it.ff]>totalnode/2)return Getcenter(it.ff,s);
  38.     }
  39.     return s;
  40. }
  41.  
  42. void dfs(int s,int p = -1,ll c = 0,int depth = 0){
  43.     lst.pb(s);
  44.     comp[s] = idx;
  45.     deep[s] = depth;
  46.     if(p!=-1){
  47.         down[s] = (down[p] * 10 + c) % M;
  48.         up[s] = c %M * power[depth-1] + up[p];
  49.         up[s]%=M;
  50.     }
  51.     for(auto it : v[s]){
  52.         if(it.ff==p||del[it.ff])continue;
  53.         if(p==-1)idx++;
  54.         dfs(it.ff,s,it.ss,depth+1);
  55.     }
  56. }
  57.  
  58. void Clear(int s){
  59.     up[s] = 0;
  60.     down[s] = 0;
  61.     idx = 0;
  62.     lst.clear();
  63.     deep[s] = 0;
  64.     dfs(s);
  65. }
  66.  
  67. ll solve(int s,int p = -1 ,ll c = 0,int depth = 0){
  68.     map<ll,ll>cnt;
  69.     map<ll,ll>cmp[idx+50];
  70.     for(auto it : lst){
  71.         cnt[up[it]]++;
  72.         cmp[comp[it]][up[it]]++;
  73.     }
  74.  
  75.     ll ret = cnt[0] - 1;
  76.  
  77.     for(auto it : lst){
  78.         if(it==s)continue;
  79.         ll val = -down[it] %M * inv[deep[it]] % M + M;
  80.         val %= M;
  81.         ret+=cnt[val];
  82.         ret-=cmp[comp[it]][val];
  83.     }
  84.     return ret;
  85. }
  86.  
  87. void Decompose(int x,int p)
  88. {
  89.     totalnode = 0;
  90.     dfs2(x,p);
  91.     int s = Getcenter(x,p);
  92.     del[s] = 1;
  93.     Clear(s);
  94.     ans+=solve(s);
  95.     ///cout<<ans<<endl;
  96.     for(auto it : v[s])
  97.     {
  98.         if(it.ff==p||del[it.ff])continue;
  99.         Decompose(it.ff,s);
  100.  
  101.     }
  102. }
  103.  
  104. int main()
  105. {
  106.     booster()
  107.     ///read("input.txt");
  108.     cin>>n>>M;
  109.     precal();
  110.     for(int i = 0;i<n-1;i++)
  111.     {
  112.         ll a,b,w;
  113.         cin>>a>>b>>w;
  114.         v[a].pb(pll(b,w));
  115.         v[b].pb(pll(a,w));
  116.     }
  117.  
  118.  
  119.     Decompose(0,-1);
  120.  
  121.     cout<<ans;
  122.  
  123.     return 0;
  124. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement