hdu 5877 Weak Pair（Treap平衡树 + dfs）

Posted 2020-08-09

题目链接：http://acm.hdu.edu.cn/showproblem.php?pid=5877

题意：在一棵树上，每个节点有一个数值，问有多少对节点满足条件。

对于一对节点(u, v)， 若满足 u 为 v的祖先节点且u，v节点上的数值乘积小于 k。

 

思路：

dfs遍历树，对于每个节点 i 查询其祖先节点中，有多少个节点的值小于等于 k / ai（ai为当前节点的数值）。

dfs过程中，动态把祖先节点插入treap树中，退栈时删除。

 

  1 #include <iostream>
  2 #include <stdio.h>
  3 #include <cstring>
  4 #include <queue>
  5 #include <set>
  6 #include <cmath>
  7 #include <time.h>
  8 #include <cstdlib>
  9 #include <algorithm>
 10 #define lson (i<<1)
 11 #define rson (lson|1)
 12 using namespace std;
 13 
 14 typedef long long LL;
 15 const int N = 100007;
 16 const int INF = 0x7fffffff;
 17 
 18 struct Node
 19 {
 20     int pri, size, ex, ch[2], p;
 21     LL val;
 22     void init()
 23     {
 24         p = 0;
 25         ch[0] = ch[1] = 0;
 26         pri = ran();
 27         val = 0;
 28         ex = 0;
 29         size = 0;
 30     }
 31     inline int ran()
 32     {
 33         int seed = 777;
 34         return (int)(seed * 48271LL % INF);
 35     }
 36 }tr[N * 20];
 37 
 38 struct Edge
 39 {
 40     int to, next;
 41 }eg[N];
 42 
 43 int n, nodeval[N], head[N], in[N], root, sum, root_treap, pos;   //pos存放节点的位置
 44 LL k, ans;
 45 
 46 void addEdge(int u, int v, int i)
 47 {
 48     eg[i].to = v;
 49     eg[i].next = head[u];
 50     head[u] = i;
 51 }
 52 
 53 void rot(int i, int ro)   //ro : 0left  1right
 54 {
 55     int son = tr[i].ch[!ro];
 56     tr[i].ch[!ro] = tr[son].ch[ro];
 57     tr[tr[son].ch[ro]].p = i;
 58     tr[son].p = tr[i].p;
 59     if (tr[son].p == 0)
 60         root_treap = son;
 61     if (root_treap != son)
 62     {
 63         if (tr[tr[i].p].ch[0] == i)
 64             tr[tr[i].p].ch[0] = son;
 65         else
 66             tr[tr[i].p].ch[1] = son;
 67     }
 68     tr[son].ch[ro] = i;
 69     tr[i].p = son;
 70 }
 71 
 72 void upadjust(int i)
 73 {
 74     while (tr[tr[i].p].pri > tr[i].pri && tr[i].p != 0)
 75     {
 76         if (tr[tr[i].p].ch[0] == i)
 77             rot(tr[i].p, 1);
 78         else
 79             rot(tr[i].p, 0);
 80     }
 81 }
 82 
 83 void insert(LL val, int pre, int dir, int i)
 84 {
 85     if (!tr[i].ex)
 86     {
 87         tr[i].init();
 88         tr[i].val = val;
 89         tr[i].size++;
 90         tr[i].ex = 1;
 91         if (pre != 0)
 92         {
 93             tr[pre].ch[dir] = i;
 94             tr[i].p = pre;
 95         }
 96         if (tr[i].p == 0)
 97             root_treap = i;
 98         else if (tr[i].pri < tr[tr[i].p].pri)
 99             upadjust(i);
100         return;
101     }
102     else if (val <= tr[i].val)
103     {
104         tr[i].size++;
105         if (tr[i].ch[0])
106             insert(val, i, 0, tr[i].ch[0]);
107         else
108             insert(val, i, 0, pos++);                     //节点按照插入treap树的顺序存储，模拟链表。
109     }
110     else
111     {
112         tr[i].size++;
113         if (tr[i].ch[1])
114             insert(val, i, 1, tr[i].ch[1]);
115         else
116             insert(val, i, 1, pos++);                     //
117     }
118 }
119 
120 void del(LL val, int flag, int i)
121 {
122     if (tr[i].val == val)
123     {
124         if (!tr[i].ch[0] && !tr[i].ch[1])
125         {
126             int dir = 0;
127             if (tr[tr[i].p].ch[1] == i)
128                 dir = 1;
129             tr[tr[i].p].ch[dir] = 0;
130             tr[i].p = 0;
131             tr[i].ex = 0;
132             tr[i].size--;
133         }
134         else if (tr[i].ch[0] && tr[i].ch[1])
135         {
136             while (tr[i].ch[0] && tr[i].ch[1])
137             {
138                 if (tr[tr[i].ch[0]].pri > tr[tr[i].ch[1]].pri)
139                     rot(i, 0);
140                 else
141                     rot(i, 1);
142             }
143             del(val, 1, root_treap);  //
144         }
145         else
146         {
147             int child = 0;
148             if (tr[i].ch[1])
149                 child = 1;
150             if (tr[tr[i].p].ch[0] == i)
151             {
152                 tr[tr[i].p].ch[0] = tr[i].ch[child];
153                 tr[tr[i].ch[child]].p = tr[i].p;
154                 if (tr[tr[i].ch[child]].p == 0)
155                     root_treap = tr[i].ch[child];
156             }
157             else
158             {
159                 tr[tr[i].p].ch[1] = tr[i].ch[child];
160                 tr[tr[i].ch[child]].p = tr[i].p;
161                 if (tr[tr[i].ch[child]].p == 0)
162                     root_treap = tr[i].ch[child];
163             }
164             tr[i].ex = 0;
165             tr[i].size--;
166         }
167     }
168     else if (tr[i].val < val)
169     {
170         if (flag == 0)
171             tr[i].size--;
172         del(val, flag, tr[i].ch[1]);
173     }
174     else
175     {
176         if (flag == 0)
177             tr[i].size--;
178         del(val, flag, tr[i].ch[0]);
179     }
180 }
181 
182 void query(LL val, int i)
183 {
184     if (tr[i].ex == 0)
185         return;
186     if (val == tr[i].val)
187     {
188         sum += (tr[tr[i].ch[0]].size + 1);
189         return;
190     }
191     else if (val > tr[i].val)
192     {
193         sum += (tr[tr[i].ch[0]].size + 1);
194         query(val, tr[i].ch[1]);
195     }
196     else
197     {
198         query(val, tr[i].ch[0]);
199     }
200 }
201 
202 void dfs(int u)
203 {
204     LL val = nodeval[u];       
205     sum = 0;
206     if (u != root)
207         query((k / val), root_treap);
208     ans += sum;
209     insert(val, 0, 0, root_treap);
210     for (int i = head[u]; i != -1; i = eg[i].next)
211     {
212         int v = eg[i].to;
213         dfs(v);
214     }
215     del(val, 0, root_treap);
216 }
217 
218 int main()
219 {
220     //freopen("in.ads", "r", stdin);
221     ios_base::sync_with_stdio(0); cin.tie(0);
222     int t;
223     while (cin >> t)
224     {
225         while (t--)
226         {
227             pos = 1;
228             ans = 0;
229             memset(in, 0, sizeof(in));
230             memset(head, -1, sizeof(head));
231             cin >> n >> k;
232             for (int i = 1; i <= n; i++)
233                 cin >> nodeval[i];
234             for (int i = 0; i < n - 1; i++)
235             {
236                 int u, v;
237                 cin >> u >> v;
238                 addEdge(u, v, i);
239                 in[v]++;
240             }
241             for (int i = 1; i <= n; i++)
242             {
243                 if (!in[i])
244                 {
245                     root = i;
246                     break;
247                 }
248             }
249             root_treap = 1;
250             dfs(root);
251             printf("%lld\n", ans);
252         }
253     }
254     return 0;
255 }