[CSP 2019]树的重心

啊...好麻烦...不想写...有空再来填吧...

  1 #include <iostream>
  2 #include <cstdio>
  3 #include <cstring>
  4 #include <algorithm>
  5 #include <cstdlib>
  6 #include <vector>
  7 #include <cctype>
  8 #define inf 300000
  9 #define ll long long
 10 #define INF 0x7fffffff
 11 
 12 namespace chiaro{
 13 
 14 template <class T>
 15 inline void read(T& num) {
 16     num = 0; register char c = getchar(), up = c;
 17     while(!isdigit(c)) up = c, c = getchar();
 18     while(isdigit(c)) num = (num << 1) + (num << 3) + (c ^ '0'), c = getchar();
 19     up == '-' ? num = -num : 0;
 20 }
 21 template <class T>
 22 inline void read(T& a,T& b) {read(a); read(b);}
 23 
 24 class bitTree {
 25 #define int long long
 26     private:
 27     int n;
 28     int c[inf];
 29     inline int lowbit(const int& x) {return x & -x;}
 30     
 31     public:
 32     inline void build(int size) {
 33         n = size; std::fill(c + 1, c + 2 + n, 0);
 34     }
 35     inline void add(int x, int k) {
 36         ++x; while(x <= n + 1) {c[x] += k; x += lowbit(x);}
 37     }
 38     inline int ask(int x) {
 39         ++x; register int res = 0; while(x) {res += c[x]; x -= lowbit(x);}
 40         return res;
 41     }
 42     inline int query(int x, int y) {
 43         return ask(x) - ask(y);
 44     }
 45 #undef int
 46 };
 47 
 48 struct edge {
 49     int to;
 50     edge* nxt;
 51 };
 52 
 53 int n, root;
 54 ll ans;
 55 int maxp, nxtmaxp;
 56 int s[inf], p[inf];
 57 bool vis[inf];
 58 edge* g[inf];
 59 bitTree c[2];
 60 
 61 bool edgeClearFlag;
 62 
 63 inline void connect(int from, int to) {
 64     static edge pool[inf << 1];
 65     static edge* p = pool;
 66     (edgeClearFlag) ? p = pool, edgeClearFlag = 0 : 0;
 67     p->to = to;
 68     p->nxt = g[from];
 69     g[from] = p; p++;
 70     return;
 71 }
 72 
 73 void dfsRoot (int now, int fa) {
 74     s[now] = 1, p[now] = 0;
 75     register bool isRoot = 1;
 76     for(register auto e = g[now]; e; e = e->nxt) {
 77         if(e->to == fa) continue;
 78         dfsRoot (e->to, now);
 79         s[now] += s[e->to];
 80         p[now] = std::max (p[now], s[e->to]);
 81         (s[e->to] << 1) > n ? isRoot = 0 : 0;
 82     }
 83     ((n - s[now]) << 1) > n ? isRoot = 0 : 0;
 84     isRoot ? root = now : 0;
 85     return;
 86 }
 87 
 88 void dfs(int now, int fa) {
 89     fa > 0 ? c[0].add(s[fa], -1), 0 : 0;
 90     c[0].add(n - s[now], 1);
 91     if(now != root) {
 92         ans += now * c[0].query(n - (p[now] << 1), n - (s[now] << 1) - 1);
 93         ans += now * c[1].query(n - (p[now] << 1), n - (s[now] << 1) - 1);
 94         (vis[now] == 0 && vis[fa]) ? vis[now] = 1 : 0;
 95         ans += root * (s[now] <= n - (s[vis[now] ? nxtmaxp : maxp] << 1));
 96     }
 97     c[1].add(s[now], 1);
 98     for(register auto e = g[now]; e; e = e->nxt) {
 99         if(e->to == fa) continue;
100         dfs(e->to, now);
101     }
102     fa > 0 ? c[0].add(s[fa], 1), 0 : 0;
103     c[0].add(n - s[now], -1);
104     (now != root) ? (ans -= now * c[1].query(n - (p[now] << 1), n - (s[now] << 1) - 1)), 0 : 0;
105 }
106 
107 inline void clear (int n) {
108     c[0].build(n); c[1].build(n);
109     std::fill(g, g + 1 + n, nullptr);
110     ans = 0; root = 0;
111     maxp = nxtmaxp = 0;
112     edgeClearFlag = 1;
113 }
114 
115 inline void solve() {
116     read(n);
117     clear(n);
118     for(register int i = 1; i < n; i++) {
119         int u, v; read(u, v);
120         connect(u, v);
121         connect(v, u);
122     }
123     dfsRoot(1, -1);
124     dfsRoot(root, -1);
125     for(register auto e = g[root]; e; e = e->nxt) {
126         (s[e->to] > s[nxtmaxp]) ? nxtmaxp = e->to : 0;
127         (s[nxtmaxp] > s[maxp]) ? std::swap(maxp, nxtmaxp), 0 : 0;
128     }
129     for(register int i = 1; i <= n; i++) c[0].add(s[i], 1);
130     std::fill(vis, vis + 1 + n, 0); vis[maxp] = 1;
131     dfs(root, -1);
132     printf("%lld\n", ans);
133     return;
134 }
135 
136 inline void setting() {
137 #ifdef ONLINE_JUDGE
138     freopen("centroid.in", "r", stdin);
139     freopen("centroid.out", "w", stdout);
140 #endif
141 }
142 
143 inline signed main () {
144     setting();
145     int T; read(T);
146     while(T--) solve();
147     return 0;
148 }
149 
150 }
151 
152 signed main() {return chiaro::main();}

 

posted @ 2020-08-27 17:14  Chiaro  阅读(151)  评论(0)    收藏  举报