【线段树 扫描线 二维数点】loj#6276. 果树

路径计数转成二维数点很妙啊

题目描述

NiroBC 姐姐是个活泼的少女，她十分喜欢爬树，而她家门口正好有一棵果树，正好满足了她爬树的需求。

这颗果树有 $N$ 个节点，标号 $1 \ldots N$ 。每个节点长着一个果子，第 $i$ 个节点上的果子颜色为 $C_i$​ 。

NiroBC 姐姐每天都要爬树，每天都要选择一条有趣的路径 $(u, v)$ 来爬。

一条路径被称作有趣的，当且仅当这条路径上的果子的颜色互不相同。

$(u, v)$ 和 $(v, u)$ 被视作同一条路径。特殊地， $(i, i)$ 也被视作一条路径，这条路径只含 $i$ 一个节点，显然是有趣的。

NiroBC 姐姐想知道这颗树上有多少条有趣的路径。

输入格式

第一行，一个整数 $N$ ，表示果树的节点数目。

接下来一行 $N$ 个整数 $C_{1 \ldots N}$ ，表示 $N$ 个果子各自的颜色。

再接下来 $N-1$ 行，每行两个整数 $u_i, v_i$​ ，表示 $u_i$​ 和 $v_i$​ 之间有一条边。

数据保证这 $N-1$ 条边构成一棵树。

输出格式

一个整数，有趣的路径的数量。

样例

样例输入 1

3
1 2 3
1 2
1 3

样例输出 1

6

样例输入 2

5
1 1 2 3 3
1 2
1 3
2 4
2 5

样例输出 2

8

样例解释 1

有 $(1,1)(1,2)(1,3)(2,2)(2,3)(3,3)$ 共 $6$ 条路径。

样例解释 2

有 $(1,1)(1,3)(2,2)(2,4)(2,5)(3,3)(4,4)(5,5)$ 共 $8$ 条路径。

数据范围与提示

对于所有数据， $1 \le N \le 10^5$ ，$1 \le C_i \le N$ ，每种颜色在树上出现不超过 $20$ 次。

本题采用打包测试。

各个 Subtask 的特殊限制如下，不填代表该项无特殊限制。

Subtask 编号	$N$	其他限制	该 Subtask 分值
0	$\le 100$		12
1	$\le 3000$		25
2		整棵树形成一条依次为 $1, 2, 3, \ldots , N$ 的链	30
3			33

---

题目分析

二维数点

注意到每种颜色的出现次数非常小，于是考虑枚举同种颜色。

对于颜色相同的一对点，它们会把树分成三个部分，而两端部分不能够相互连边。这个部分用dfs序把图再构一遍，就方便处理了。至于两种分别是成链不成链的情况，稍微细心点细节就没问题。

点对$(x,y)$可以表示为二维点$(x,y)$，也就是说每一次将非法位置转为二维矩形覆盖，那么最后就是枚举每一个点，对非法矩形差分扫描线处理。

码力还是不够，标记${\rm //HERE}$都是打挂的地方。

 #include<bits/stdc++.h>
 const int maxn = ;
 const int maxm = ;
 const int maxLog = ;

 struct dfn
 {
     int l,r;
 }a[maxn];
 struct node
 {
     int mn,tot,tag;
 }f[maxn<<];
 struct line
 {
     int l,r,opt;
     line (int a=, int b=, int c=):l(a),r(b),opt(c) {}
 };
 long long ans;
 int n,c[maxn],fa[maxn][maxLog],dep[maxn],chTot;
 int edgeTot,head[maxn],nxt[maxm],edges[maxm];
 std::vector<int> col[maxn];
 std::vector<line> mdf[maxn];

 int read()
 {
     char ch = getchar();
     int num = , fl = ;
     for (; !isdigit(ch); ch=getchar())
         if (ch=='-') fl = -;
     for (; isdigit(ch); ch=getchar())
         num = (num<<)+(num<<)+ch-;
     return num*fl;
 }
 void addedge(int u, int v)
 {
     edges[++edgeTot] = v, nxt[edgeTot] = head[u], head[u] = edgeTot;
     edges[++edgeTot] = u, nxt[edgeTot] = head[v], head[v] = edgeTot;
 }
 inline int kfa(int x, int d)
 {
     if (d==-) return x;
     for (int i=; i>=; i--)
         if (d>>i&) x = fa[x][i];
     return x;
 }
 void addLine(int l1, int r1, int l2, int r2)
 {
     if (l1 > r1||l2 > r2) return;
     mdf[l1].push_back(line(l2, r2, ));
     mdf[r1+].push_back(line(l2, r2, -));
     mdf[l2].push_back(line(l1, r1, ));
     mdf[r2+].push_back(line(l1, r1, -));
 }
 inline void split(int x, int y)
 {
     if (dep[x] > dep[y]) std::swap(x, y);
     int anc = kfa(y, dep[y]-dep[x]-);        //HERE
     if (fa[anc][]!=x||(dep[x]==dep[y]&&x!=y)) addLine(a[x].l, a[x].r, a[y].l, a[y].r);
     else{
         addLine(, a[anc].l-, a[y].l, a[y].r);
         addLine(a[anc].r+, n, a[y].l, a[y].r);        //HERE　　　　这里打挂只剩链30pts
     }
 }
 void smcolConnect()
 {
     register int i,j,k;
     for (i=; i<=n; i++)
         for (j=; j<col[i].size(); j++)
             for (k=j+; k<col[i].size(); k++)
                 split(col[i][j], col[i][k]);        //HERE
 }
 void dfs(int x, int fat)
 {
     dep[x] = dep[fat]+, fa[x][] = fat;
     a[x].l = ++chTot;
     for (int i=head[x]; i!=-; i=nxt[i])
         if (edges[i]!=fat) dfs(edges[i], x);
     a[x].r = chTot;
 }
 void build(int rt, int l, int r)
 {
     f[rt].tot = r-l+;
     if (l==r) return;
     int mid = (l+r)>>;
     build(rt<<, l, mid);
     build(rt<<|, mid+, r);
 }
 void init()
 {
     for (int j=; j<=; j++)
         for (int i=; i<=n; i++)
             fa[i][j] = fa[fa[i][j-]][j-];
     build(, , n);
 }
 void pushdown(int rt)
 {
     int l = rt<<, r = rt<<|, &t = f[rt].tag;
     if (t){
         f[l].tag += t, f[r].tag += t;
         f[l].mn += t, f[r].mn += t, t = ;            //HERE
     }
 }
 void pushup(int rt)
 {
     f[rt].mn = std::min(f[rt<<].mn, f[rt<<|].mn);
     f[rt].tot = (f[rt].mn==f[rt<<].mn?f[rt<<].tot:)+(f[rt].mn==f[rt<<|].mn?f[rt<<|].tot:);
 }
 void modify(int rt, int L, int R, int l, int r, int c)
 {
     if (L <= l&&r <= R){
         f[rt].mn += c, f[rt].tag += c;
         return;
     }
     int mid = (l+r)>>;
     pushdown(rt);　　　　　　//HERE
     if (L <= mid) modify(rt<<, L, R, l, mid, c);
     if (R > mid) modify(rt<<|, L, R, mid+, r, c);
     pushup(rt);
 }
 int main()
 {
     memset(head, -, sizeof head);
     n = read();
     for (int i=; i<=n; i++) c[i] = read(), col[c[i]].push_back(i);
     for (int i=; i<n; i++) addedge(read(), read());
     dfs(, ), init();
     smcolConnect();
     for (int i=; i<=n; i++)
     {
         for (unsigned int p=; p<mdf[i].size(); p++)
             modify(, mdf[i][p].l, mdf[i][p].r, , n, mdf[i][p].opt);
         if (f[].mn==) ans += f[].tot;
     }
     printf("%lld\n",(ans+n)>>);
     return ;
 }

DSU做法

还看到一种神奇的dsu+可持久化线段树做法：LibreOJ #6276.果树 dsu on tree+可持久化线段树

END