#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
#define rep(i,s,t) for(register ll i = s;i <= t;++i)
#define per(i,t,s) for(register ll i = t;i >= s;--i)
const ll N = 1e6 + 5;
ll n;
ll k;
ll rt1;
ll rt2;
ll top;
ll idx;
ll ans;
ll p[N] = {};
ll q[N] = {};
ll fa[N] = {};
ll st[N] = {};
ll sz[N] = {};
ll siz[N] = {};
ll dfn[N] = {};
ll son[N] = {};
vector<ll> g[N];
vector<ll> g1[N];
vector<ll> g2[N];
class binary_indexed_tree
{
private:
ll t[N] = {};
public:
inline void init()
{
memset(t,0,sizeof(t));
}
inline ll lowbit(ll x)
{
return x & (-x);
}
inline void upd(ll x,ll k)
{
while(x <= n)
{
t[x] += k;
x += lowbit(x);
}
}
inline ll qry(ll x)
{
ll ans = 0;
while(x)
{
ans += t[x];
x -= lowbit(x);
}
return ans;
}
};
binary_indexed_tree t1;
binary_indexed_tree t2;
inline ll read()
{
ll x = 0;
ll y = 1;
char c = getchar();
while(c < '0' || c > '9')
{
if(c == '-')
y = -y;
c = getchar();
}
while(c >= '0' && c <= '9')
{
x = (x << 3) + (x << 1) + (c ^ '0');
c = getchar();
}
return x * y;
}
inline void write(ll x)
{
if(x < 0)
{
putchar('-');
write(-x);
return;
}
if(x > 9)
write(x / 10);
putchar(x % 10 + '0');
}
inline void dfs(ll u)
{
siz[u] = 1;
dfn[u] = ++idx;
for(register auto v : g1[u])
{
dfs(v);
siz[u] += siz[v];
}
}
inline void dfs1(ll u)
{
st[++top] = u;
g[u].clear();
if(top > k)
fa[u] = st[top - k];
else
fa[u] = 0;
if(fa[u])
g[fa[u]].push_back(u);
sz[u] = 1;
son[u] = 0;
for(auto v : g2[u])
{
dfs1(v);
if(sz[v] > sz[son[u]])
son[u] = v;
sz[u] += sz[v];
}
top--;
}
inline void ins(ll x,ll k)
{
t1.upd(dfn[x],k);
t1.upd(dfn[x] + siz[x],-k);
t2.upd(dfn[x],k);
}
inline ll query(ll x)
{
return t1.qry(dfn[x]) + t2.qry(dfn[x] + siz[x] - 1) - t2.qry(dfn[x] - 1);
}
inline void dfs3(ll u,ll k)
{
for(auto v : g[u])
{
if(k == 1)
ans += query(v);
else if(k == 2)
ins(v,1);
else if(k == 3)
ins(v,-1);
}
for(auto v : g2[u])
dfs3(v,k);
}
inline void dfs2(ll u,ll k) // k is keep flag
{
for(auto v: g2[u])
if(v != son[u])
dfs2(v,0);
if(son[u])
dfs2(son[u],1);
for(auto v: g2[u])
if(v != son[u])
{
dfs3(v,1);
dfs3(v,2);
}
for(register auto v : g[u])
ins(v,1);
if(!k)
dfs3(u,3);
}
// 添加重置全局状态的函数
inline void reset_global() {
// 重置 DFS 相关全局状态
top = 0;
idx = 0;
// 重置树状数组在 cal 逻辑中处理,不在此重置
// 重置 DSU 相关数组
memset(sz, 0, sizeof(sz));
memset(son, 0, sizeof(son));
memset(fa, 0, sizeof(fa));
memset(st, 0, sizeof(st));
// g 数组在 dfs1 中每个节点清空,无需全局重置
}
// 封装 cal 函数,与第一段代码一致
inline void cal() {
reset_global(); // 重置全局状态
dfs(rt1); // 在 T1 上 DFS
t1.init(); // 清空树状数组
t2.init();
top = 0;
dfs1(rt2); // 在 T2 上处理第 k 祖先
dfs2(rt2, 0); // DSU on tree
}
int main()
{
freopen("D.in","r",stdin);
freopen("D.out","w",stdout);
n = read();
k = read();
rep(i,1,n)
p[i] = read();
rep(i,1,n)
q[i] = read();
rep(i,1,n)
{
if(!p[i])
rt1 = i;
else
g1[p[i]].push_back(i);
if(!q[i])
rt2 = i;
else
g2[q[i]].push_back(i);
}
// 第一次计算
cal();
// 交换树
rep(i,1,n)
{
swap(p[i],q[i]);
swap(g1[i],g2[i]);
}
swap(rt1,rt2);
// 第二次计算
cal();
write(ans);
fclose(stdin);
fclose(stdout);
return 0;
}小丁的树
题目描述
小丁拥有两棵均具有
n
n 个顶点,编号集合为
{
1
,
2
,
⋯
,
n
}
{1,2,⋯,n} 的有根树
T
1
,
T
2
T
1
,T
2
,现在他需要计算这两棵树的相似程度。
为了计算,小丁定义了对于一棵树
T
T 和
T
T 上两个不同顶点
u
,
v
u,v 的距离函数
d
T
(
u
,
v
)
d
T
(u,v),其定义为
u
,
v
u,v 两个点距离成为祖先关系有多近,具体来说,对于所有在
T
T 上为祖先关系的点对
(
u
′
,
v
′
)
(u
′
,v
′
),
dis
(
u
,
u
′
)
+
dis
(
v
,
v
′
)
dis(u,u
′
)+dis(v,v
′
) 的最小值即为
d
T
(
u
,
v
)
d
T
(u,v) 的值,其中
dis
(
u
,
v
)
dis(u,v) 表示
u
,
v
u,v 在树
T
T 上的唯一简单路径包含的边数,即
u
,
v
u,v 的距离。
点对
(
u
′
,
v
′
)
(u
′
,v
′
) 为祖先关系,当且仅当
u
′
u
′
是
v
′
v
′
的祖先或
v
′
v
′
是
u
′
u
′
的祖先。(注意,每个点都是自己的祖先)
小丁心里还有一个参数
k
k,如果节点对
(
u
,
v
)
(u,v) 满足以下条件,称之为不相似的节点对:
1
≤
u
<
v
≤
n
1≤u<v≤n
"
d
T
1
(
u
,
v
)
=
0
d
T
1
(u,v)=0 且
d
T
2
(
u
,
v
)
>
k
d
T
2
(u,v)>k“ 或 "
d
T
2
(
u
,
v
)
=
0
d
T
2
(u,v)=0 且
d
T
1
(
u
,
v
)
>
k
d
T
1
(u,v)>k“
小丁认为,不相似的节点对越多,
T
1
T
1
和
T
2
T
2
就越不相似,你能告诉他总共有多少不相似的节点对吗?
输入格式
第一行两个整数
n
,
k
n,k,表示
T
1
T
1
和
T
2
T
2
的节点数和参数
k
k。
第二行
n
n 个正整数
p
1
,
p
2
,
⋯
,
p
n
p
1
,p
2
,⋯,p
n
,
T
1
T
1
中节点
i
i 的父节点为
p
i
p
i
,特别的,若
p
i
=
0
p
i
=0,则
i
i 是
T
1
T
1
的根。
第三行
n
n 个正整数
q
1
,
q
2
,
⋯
,
q
n
q
1
,q
2
,⋯,q
n
,
T
2
T
2
中节点
i
i 的父节点为
q
i
q
i
,特别的,若
q
i
=
0
q
i
=0,则
i
i 是
T
2
T
2
的根。
输出格式
一行一个整数,表示不相似的节点对总数。
样例 1 输入
5 0
0 1 1 2 3
5 3 1 1 0
样例 1 输出
4
样例 1 解释
(
2
,
3
)
,
(
2
,
4
)
,
(
2
,
5
)
,
(
4
,
5
)
(2,3),(2,4),(2,5),(4,5) 为不相似的节点对。
其余样例见下发文件。
数据规模与约定
对于所有数据,
1
≤
n
≤
2
×
10
5
,
0
≤
k
<
n
,
0
≤
p
i
,
q
i
≤
n
1≤n≤2×10
5
,0≤k<n,0≤p
i
,q
i
≤n,且由
p
i
,
q
i
p
i
,q
i
形成的是一棵
n
n 个节点的有根树。
本题采用捆绑评测,你只有通过了一个子任务中所有测试点才能得到该子任务的分数。
Subtask 1(10pts):
1
≤
n
≤
100
1≤n≤100。
Subtask 2(20pts):
1
≤
n
≤
3000
1≤n≤3000。
Subtask 3(20pts):
k
=
0
k=0。
Subtask 4(10pts):
0
≤
k
≤
20
0≤k≤20。
Subtask 5(40pts):无特殊限制。
请详细注释上述代码的每一行,给出详细的解题思路,并完整按照代码流程模拟样例