fork download
copy 
  1. /**
  2.  *   /\_/\
  3.  *  (= ._.)
  4.  *  / >01<\
  5.  *
  6. */
  7.  
  8. #include<bits/stdc++.h>
  9. #define task ""
  10. #define int long long
  11. #define ii pair<int, int>
  12. #define ll long long
  13. #define ld long double
  14. #define fi first
  15. #define se second
  16. #define faster ios_base::sync_with_stdio(NULL); cin.tie(NULL);
  17. #define F if(fopen(task".inp","r")){ freopen(task".inp","r",stdin); freopen(task".out","w",stdout); }
  18. using namespace std;
  19.  
  20. const int Q = 1e6 + 21;
  21. const int N = 1e3 + 21;
  22. const int LOG = 20;
  23. const int mod = 1e9 + 7;
  24. const int inf = 1e9 + 7;
  25. const int dx[] = {0, 0, 1, -1};
  26. const int dy[] = {1, -1, 0, 0};
  27.  
  28. int n, t;
  29. int a[Q];
  30.  
  31. vector<int> adj[Q];
  32.  
  33. struct pt
  34. {
  35.     int u, v, w;
  36.  
  37.     bool operator<(pt other) const {
  38.         return w < other.w;
  39.     }
  40. };
  41.  
  42. struct DSU
  43. {
  44.     vector<int> par;
  45.  
  46.     DSU (int n)
  47.     {
  48.         par.resize(n + 6, -1);
  49.     }
  50.  
  51.     int fpar(int u)
  52.     {
  53.         if(par[u] < 0)
  54.             return u;
  55.  
  56.         return par[u] = fpar(par[u]);
  57.     }
  58.  
  59.     bool join(int u, int v)
  60.     {
  61.         u = fpar(u);
  62.         v = fpar(v);
  63.  
  64.         if(u == v)
  65.             return false;
  66.  
  67.         if(par[u] > par[v])
  68.             swap(u, v);
  69.  
  70.         par[u] += par[v];
  71.         par[v] = u;
  72.  
  73.         return true;
  74.     }
  75.  
  76.     int get_sz(int u)
  77.     {
  78.         return -par[fpar(u)];
  79.     }
  80. };
  81.  
  82. int power(int a, int b)
  83. {
  84.     int res = 1;
  85.  
  86.     while(b)
  87.     {
  88.         if(b & 1)
  89.             res = (res * a) % mod;
  90.  
  91.         a = (a * a) % mod;
  92.         b >>= 1;
  93.     }
  94.  
  95.     return res;
  96. }
  97.  
  98. int fac[Q], inv[Q];
  99.  
  100. void init()
  101. {
  102.     fac[0] = 1;
  103.  
  104.     for(int i = 1; i < Q; i++)
  105.         fac[i] = (fac[i - 1] * i) % mod;
  106.  
  107.     inv[Q - 1] = power(fac[Q - 1], mod - 2);
  108.  
  109.     for(int i = Q - 2; i >= 0; i--)
  110.         inv[i] = (inv[i + 1] * (i + 1)) % mod;
  111. }
  112.  
  113. int nCk(int n, int k)
  114. {
  115.     return fac[n] * inv[k] % mod * inv[n - k] % mod;
  116. }
  117.  
  118. int sz[Q];
  119.  
  120. int dfs(int u, int par, int S)
  121. {
  122.     sz[u] = 1;
  123.     int cur_sum = 0;
  124.     int all = nCk(S, 3);
  125.  
  126.     for(int v : adj[u])
  127.     {
  128.         if(v == par)
  129.             continue;
  130.  
  131.         cur_sum = (cur_sum + dfs(v, u, S)) % mod;
  132.         sz[u] += sz[v];
  133.  
  134.         int s1 = sz[v];
  135.         int s2 = S - sz[v];
  136.  
  137.         int ways = (all - nCk(s1, 3) - nCk(s2, 3)) % mod;
  138.  
  139.         if(ways < 0)
  140.             ways += mod;
  141.  
  142.         cur_sum = (cur_sum + ways) % mod;
  143.     }
  144.  
  145.     return cur_sum;
  146. }
  147.  
  148. bool ce[Q];
  149.  
  150. int cal(int t, DSU dsu)
  151. {
  152.     memset(ce, false, sizeof(ce));
  153.  
  154.     int S = 0;
  155.     int ans = 0;
  156.  
  157.     for(int u = 1; u <= n; u++)
  158.     {
  159.         int pu = dsu.fpar(u);
  160.  
  161.         if(ce[pu])
  162.             continue;
  163.  
  164.         ce[pu] = true;
  165.         S = dsu.get_sz(pu);
  166.  
  167.         if(S >= 3)
  168.         {
  169.             int sum = dfs(u, 0, S);
  170.             ans = (ans + sum) % mod;
  171.         }
  172.     }
  173.  
  174.     return (2 * ans) % mod;
  175. }
  176.  
  177. void solve()
  178. {
  179.     cin >> n;
  180.     cin >> t;
  181.  
  182.     init();
  183.  
  184.     for(int i = 1; i <= n; i++)
  185.         cin >> a[i];
  186.  
  187.     vector<pt> edge;
  188.  
  189.     for(int i = 1; i < n; i++)
  190.     {
  191.         int u, v;
  192.  
  193.         cin >> u;
  194.         cin >> v;
  195.  
  196.         edge.push_back({u, v, max(a[u], a[v])});
  197.     }
  198.  
  199.     sort(edge.begin(), edge.end());
  200.  
  201.     DSU dsu (n);
  202.  
  203.     for(auto [u, v, w] : edge)
  204.     {
  205.         if(w >= t)
  206.             break;
  207.  
  208.         dsu.join(u, v);
  209.  
  210.         adj[u].push_back(v);
  211.         adj[v].push_back(u);
  212.     }
  213.  
  214.     int t1 = cal(t - 1, dsu);
  215.  
  216.     for(auto [u, v, w] : edge)
  217.     {
  218.         if(w > t)
  219.             break;
  220.  
  221.         if(w != t)
  222.             continue;
  223.  
  224.         dsu.join(u, v);
  225.  
  226.         adj[u].push_back(v);
  227.         adj[v].push_back(u);
  228.     }
  229.  
  230.     int t2 = cal(t, dsu);
  231.  
  232.     cout << (t2 - t1 + mod) % mod;
  233.     cout << '\n';
  234. }
  235.  
  236. signed main()
  237. {
  238.     faster;
  239.     F;
  240.  
  241.     int test = 1;
  242.  
  243. //    cin >> test;
  244.  
  245.     while(test--)
  246.     {
  247.         solve();
  248.     }
  249.  
  250.     cerr << "Time elapsed: " << (1.0 * clock() / CLOCKS_PER_SEC) << " s.\n";
  251.  
  252.     return 0;
  253. }
  254.  
  255.  
Success #stdin #stdout #stderr 0.03s 47224KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
0
stderr
copy 
Time elapsed: 0.0238 s.
https://ideone.com/LplgD9
language:
C++ (gcc 8.3)
created:
8 days ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!