#P9603. [IOI2023] 山毛榉树

[IOI2023] 山毛榉树

题目背景

这是一道交互题,你只需要实现代码中要求的函数。

你的代码不需要引用任何额外的头文件,也不需要实现 main 函数。

本题仅支持 C++ 语言提交。

题目描述

Vétyem Woods 是一片著名的缤纷多彩的森林。其中最老最高的一棵山毛榉树叫 Ős Vezér。

树 Ős Vezér 可以被建模成 NN结点N1N-1的集合。结点的编号为从 00N1N-1,边的编号为从 11N1N-1。每条边均连接树上两个不同的结点。具体地说,边 ii1i<N1 \le i \lt N)从结点 ii 连接到结点 P[i]P[i],这里 0P[i]<i0 \le P[i] \lt i。结点 P[i]P[i] 被称为是结点 ii父结点,而结点 ii 被称为是结点 P[i]P[i] 的一个子结点

每条边都有某种颜色。一共有 MM 种可能的颜色,编号为从 11MM。边 ii 的颜色为 C[i]C[i]。不同的边可能有相同的颜色。

注意,在上面的定义中,i=0i = 0 的情形并不对应树上的边。方便起见,我们令 P[0]=1P[0] = -1C[0]=0C[0] = 0

例如,假定 Ős Vezér 有 N=18N = 18 个结点和 M=3M = 3 种可能的颜色,以及 1717 条边。边的描述为 $P = [-1, 0, 0, 0, 1, 1, 1, 2, 2, 3, 3, 3, 4, 4, 5, 10, 11, 11]$,边的颜色为 $C = [0, 1, 2, 3, 1, 2, 3, 1, 3, 3, 2, 1, 1, 2, 2, 1, 2, 3]$。这棵树如下图所示:

Árpád 是一位才华横溢的护林人,他喜欢研究树上被称为子树的部分。 对所有满足 0r<N0 \le r \lt Nrr,结点 rr 的子树是一个满足以下性质的结点集合 T(r)T(r)

  • 结点 rr 属于 T(r)T(r)
  • 如果某个结点 xx 属于 T(r)T(r),则 xx 的所有子结点都属于T(r)T(r)
  • 除了上述情况以外,其他结点都不属于 T(r)T(r)

集合 T(r)T(r) 的大小记作 T(r)|T(r)|

Árpád 最近发现了一个复杂但有趣的子树性质。Árpád 的发现需要用到大量的纸和笔做演算,他认为你需要做同样的事情才能完成理解。他还会给你几个例子,让你能够对它们做详细的分析。

假设我们有某个给定的 rr,以及子树 T(r)T(r) 中结点的某个置换 v0,v1,,vT(r)1v_0, v_1, \ldots, v_{|T(r)|-1}

对于所有满足 1i<T(r)1 \le i \lt |T(r)|ii,令 f(i)f(i) 为颜色 C[vi]C[v_i] 在长为 i1i-1 的颜色序列 C[v1],C[v2],,C[vi1]C[v_1], C[v_2], \ldots, C[v_{i-1}] 中的出现次数。

(注意,f(1)f(1) 必定为 00,原因是其定义中要考察的颜色序列是空的。)

置换 v0,v1,,vT(r)1v_0, v_1, \ldots, v_{|T(r)|-1} 被称为是一个绝妙置换,当且仅当以下性质成立:

  • v0=rv_0 = r
  • 对于所有满足 1i<T(r)1 \le i \lt |T(r)|ii,结点 viv_i 的父结点是 vf(i)v_{f(i)}

对于所有满足 0r<N0 \le r \lt Nrr,子树 T(r)T(r) 是一棵绝妙子树,当且仅当 T(r)T(r) 中结点存在某个绝妙置换。注意,根据定义,仅包含单独一个结点的子树都是绝妙的。

考虑上面给出的树的例子。可以看到,子树 T(0)T(0)T(3)T(3) 不是绝妙的。子树 T(14)T(14) 是绝妙的,因为它仅包含一个结点。接下来,我们将要说明子树 T(1)T(1) 也是绝妙的。

考虑一个由不同整数构成的序列 $[v_0, v_1, v_2, v_3, v_4, v_5, v_6] = [1, 4, 5, 12, 13, 6, 14]$。这个序列是 T(1)T(1) 中结点的一个置换。下图给出了这个置换。序列中每个结点旁边的数字,是该结点在置换中的索引。

我们将要验证,这是一个绝妙置换

  • v0=1v_0 = 1
  • f(1)=0f(1) = 0,原因是 C[v1]=C[4]=1C[v_1] = C[4] = 1 在序列 [][\,] 中出现了 00 次。
  • 相应地,v1v_1 的父结点是 v0v_0。也就是说,44 的父结点是 11。(形式化地,P[4]=1P[4] = 1。)
  • f(2)=0f(2) = 0,原因是 C[v2]=C[5]=2C[v_2] = C[5] = 2 在序列 [1][1] 中出现了 00 次。
  • 相应地,v2v_2 的父结点是 v0v_0。也就是说,55 的父结点是 11
  • f(3)=1f(3) = 1,原因是 C[v3]=C[12]=1C[v_3] = C[12] = 1 在序列 [1,2][1, 2] 中出现了 11 次。
  • 相应地,v3v_3 的父结点是 v1v_1。也就是说,1212 的父结点是 44
  • f(4)=1f(4) = 1,原因是 C[v4]=C[13]=2C[v_4] = C[13] = 2 在序列 [1,2,1][1, 2, 1] 中出现了 11 次。
  • 相应地,v4v_4 的父结点是 v1v_1。也就是说,1313 的父结点是 4。
  • f(5)=0f(5) = 0,原因是 C[v5]=C[6]=3C[v_5] = C[6] = 3 在序列 [1,2,1,2][1, 2, 1, 2] 中出现了 00 次。
  • 相应地,v5v_5 的父结点是 v0v_0。也就是说,66 的父结点是 11
  • f(6)=2f(6) = 2,原因是 C[v6]=C[14]=2C[v_6] = C[14] = 2 在序列 [1,2,1,2,3][1, 2, 1, 2, 3] 中出现了 22 次。
  • 相应地,v6v_6 的父结点是 v2v_2。也就是说,1414 的父结点是 55

由于我们能为 T(1)T(1) 中的结点找到一个绝妙置换,子树 T(1)T(1) 因此是一棵绝妙子树

你的任务是,帮助 Árpád 确定 Ős Vezér 的每棵子树是否是绝妙的。

输入格式

你需要实现以下函数。

int[] beechtree(int N, int M, int[] P, int[] C)
  • NN:树中的结点数量。
  • MM:树中边的可能颜色的数量。
  • PPCC:长度为 NN 的两个数组,以描述树中的边。
  • 该函数应当返回长度为 NN 的某个数组 bb。 对所有满足 0r<N0 \le r \lt Nrr,如果 T(r)T(r) 是绝妙的,则 b[r]b[r] 应为 11,否则应为 00
  • 该函数在每个测试用例上恰好被调用一次。

提示

样例

样例 1

考虑如下调用:

beechtree(4, 2, [-1, 0, 0, 0], [0, 1, 1, 2])

这棵树如下图所示:

T(1)T(1)T(2)T(2)T(3)T(3) 均各自包含单独一个结点,因此都是绝妙的。 T(0)T(0) 不是绝妙的。 因此,函数应当返回 [0,1,1,1][0, 1, 1, 1]

样例 2

考虑如下调用:

beechtree(18, 3, 
          [-1, 0, 0, 0, 1, 1, 1, 2, 2, 3, 3, 3, 4, 4, 5, 10, 11, 11],
          [0, 1, 2, 3, 1, 2, 3, 1, 3, 3, 2, 1, 1, 2, 2, 1, 2, 3])

这个例子在题面中已经给出。

函数应当返回 $[0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]$。

样例 3

考虑如下调用:

beechtree(7, 2, [-1, 0, 1, 1, 0, 4, 5], [0, 1, 1, 2, 2, 1, 1])

该例子如下图所示。

T(0)T(0) 是唯一不是绝妙的子树。函数应当返回 [0,1,1,1,1,1,1][0, 1, 1, 1, 1, 1, 1]

约束条件

  • 3N2000003 \le N \le 200\,000
  • 2M2000002 \le M \le 200\,000
  • 0P[i]<i0 \le P[i] \lt i(对于所有满足 1i<N1 \le i \lt Nii
  • 1C[i]M1 \le C[i] \le M(对于所有满足 1i<N1 \le i \lt Nii
  • P[0]=1P[0] = -1C[0]=0C[0] = 0

子任务

  1. (9 分)N8N \le 8M500M \le 500
  2. (5 分)边 ii 从结点 ii 连接到结点 i1i-1。也就是说,对所有满足 1i<N1 \le i \lt Nii,都有 P[i]=i1P[i] = i-1
  3. (9 分)除了结点 00 以外,其他结点要么连接到结点 00,要么连接到某个连接到结点 00 的结点。 也就是说,对于所有满足 1i<N1 \le i \lt Nii,要么有 P[i]=0P[i]=0,要么有 P[P[i]]=0P[P[i]]=0
  4. (8 分)对于所有满足 1cM1 \le c \le Mcc,至多有两条边的颜色为 cc
  5. (14 分) N200N \le 200M500M \le 500
  6. (14 分) N2000N \le 2\,000M=2M = 2
  7. (12 分) N2000N \le 2\,000
  8. (17 分) M=2M = 2
  9. (12 分) 没有额外的约束条件。

评测程序示例

评测程序示例按以下格式读取输入:

  • 11 行:N  MN \; M
  • 22 行:P[0]  P[1]    P[N1]P[0] \; P[1] \; \ldots \; P[N-1]
  • 33 行:C[0]  C[1]    C[N1]C[0] \; C[1] \; \ldots \; C[N-1]

b[0],  b[1],  b[0], \; b[1], \; \ldots 表示 beechtree 所返回的数组中的元素。评测程序示例以如下格式,在单行中输出你的答案:

  • 11 行:b[0]  b[1]  b[0] \; b[1] \; \ldots