#P9603. [IOI 2023] 山毛榉树

[IOI 2023] 山毛榉树

Background

This is an interactive problem. You only need to implement the function required in the code.

Your code must not include any additional header files, and you do not need to implement the main function.

This problem only supports C++ submissions.

Problem Description

Vétyem Woods is a famous, colorful forest. The oldest and tallest beech tree is called Ős Vezér.

The tree Ős Vezér can be modeled as a set of NN nodes and N1N-1 edges. The nodes are numbered from 00 to N1N-1, and the edges are numbered from 11 to N1N-1. Each edge connects two different nodes in the tree. Specifically, edge ii (1i<N1 \le i \lt N) connects node ii to node P[i]P[i], where 0P[i]<i0 \le P[i] \lt i. Node P[i]P[i] is called the parent of node ii, and node ii is called a child of node P[i]P[i].

Each edge has a color. There are MM possible colors, numbered from 11 to MM. The color of edge ii is C[i]C[i]. Different edges may have the same color.

Note that in the definition above, the case i=0i = 0 does not correspond to an edge in the tree. For convenience, we set P[0]=1P[0] = -1 and C[0]=0C[0] = 0.

For example, suppose Ős Vezér has N=18N = 18 nodes and M=3M = 3 possible colors, and thus 1717 edges. The edges are described by $P = [-1, 0, 0, 0, 1, 1, 1, 2, 2, 3, 3, 3, 4, 4, 5, 10, 11, 11]$, and the edge colors are $C = [0, 1, 2, 3, 1, 2, 3, 1, 3, 3, 2, 1, 1, 2, 2, 1, 2, 3]$. The tree is shown in the figure below:

Árpád is a talented forest ranger who likes to study parts of a tree called subtrees.
For each rr with 0r<N0 \le r \lt N, the subtree of node rr is a set of nodes T(r)T(r) with the following properties:

  • Node rr is in T(r)T(r).
  • If a node xx is in T(r)T(r), then all children of xx are also in T(r)T(r).
  • No other nodes are in T(r)T(r).

The size of the set T(r)T(r) is denoted by T(r)|T(r)|.

Árpád recently discovered a complex but interesting property of subtrees. His discovery requires a lot of paper-and-pencil work, and he thinks you need to do the same to fully understand it. He will also give you a few examples so that you can analyze them in detail.

Suppose we have some fixed rr, and a permutation v0,v1,,vT(r)1v_0, v_1, \ldots, v_{|T(r)|-1} of the nodes in the subtree T(r)T(r).

For each ii with 1i<T(r)1 \le i \lt |T(r)|, let f(i)f(i) be the number of occurrences of the color C[vi]C[v_i] in the color sequence of length i1i-1, namely C[v1],C[v2],,C[vi1]C[v_1], C[v_2], \ldots, C[v_{i-1}].

(Note that f(1)f(1) is always 00, because the sequence in its definition is empty.)

The permutation v0,v1,,vT(r)1v_0, v_1, \ldots, v_{|T(r)|-1} is called a nice permutation if and only if the following conditions hold:

  • v0=rv_0 = r.
  • For each ii with 1i<T(r)1 \le i \lt |T(r)|, the parent of node viv_i is vf(i)v_{f(i)}.

For each rr with 0r<N0 \le r \lt N, the subtree T(r)T(r) is called a nice subtree if and only if there exists some nice permutation of the nodes in T(r)T(r). Note that, by definition, any subtree that contains only one node is nice.

Consider the example tree given above. You can see that subtrees T(0)T(0) and T(3)T(3) are not nice. Subtree T(14)T(14) is nice because it contains only one node. Next, we will show that subtree T(1)T(1) is also nice.

Consider the sequence of distinct integers $[v_0, v_1, v_2, v_3, v_4, v_5, v_6] = [1, 4, 5, 12, 13, 6, 14]$. This sequence is a permutation of the nodes in T(1)T(1). The figure below shows this permutation. The number next to each node in the sequence is the node’s index in the permutation.

We will verify that this is a nice permutation.

  • v0=1v_0 = 1.
  • f(1)=0f(1) = 0, because C[v1]=C[4]=1C[v_1] = C[4] = 1 occurs 00 times in the sequence [][\,].
    • Accordingly, the parent of v1v_1 is v0v_0. That is, the parent of 44 is 11. (Formally, P[4]=1P[4] = 1.)
  • f(2)=0f(2) = 0, because C[v2]=C[5]=2C[v_2] = C[5] = 2 occurs 00 times in the sequence [1][1].
    • Accordingly, the parent of v2v_2 is v0v_0. That is, the parent of 55 is 11.
  • f(3)=1f(3) = 1, because C[v3]=C[12]=1C[v_3] = C[12] = 1 occurs 11 time in the sequence [1,2][1, 2].
    • Accordingly, the parent of v3v_3 is v1v_1. That is, the parent of 1212 is 44.
  • f(4)=1f(4) = 1, because C[v4]=C[13]=2C[v_4] = C[13] = 2 occurs 11 time in the sequence [1,2,1][1, 2, 1].
    • Accordingly, the parent of v4v_4 is v1v_1. That is, the parent of 1313 is 44.
  • f(5)=0f(5) = 0, because C[v5]=C[6]=3C[v_5] = C[6] = 3 occurs 00 times in the sequence [1,2,1,2][1, 2, 1, 2].
    • Accordingly, the parent of v5v_5 is v0v_0. That is, the parent of 66 is 11.
  • f(6)=2f(6) = 2, because C[v6]=C[14]=2C[v_6] = C[14] = 2 occurs 22 times in the sequence [1,2,1,2,3][1, 2, 1, 2, 3].
    • Accordingly, the parent of v6v_6 is v2v_2. That is, the parent of 1414 is 55.

Since we can find a nice permutation for the nodes in T(1)T(1), the subtree T(1)T(1) is therefore a nice subtree.

Your task is to help Árpád determine whether each subtree of Ős Vezér is nice.

Input Format

You need to implement the following function.

int[] beechtree(int N, int M, int[] P, int[] C)
  • NN: the number of nodes in the tree.
  • MM: the number of possible edge colors in the tree.
  • PP, CC: two arrays of length NN describing the edges of the tree.
  • The function should return an array bb of length NN.
    For each rr with 0r<N0 \le r \lt N, if T(r)T(r) is nice, then b[r]b[r] should be 11, otherwise it should be 00.
  • This function is called exactly once for each test case.

Hint

Samples

Sample 1

Consider the following call:

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

The tree is shown in the figure below:

T(1)T(1), T(2)T(2), and T(3)T(3) each contain a single node, so they are all nice.
T(0)T(0) is not nice.
Therefore, the function should return [0,1,1,1][0, 1, 1, 1].

Sample 2

Consider the following call:

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])

This example has already been given in the statement.

The function should return $[0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]$.

Sample 3

Consider the following call:

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

This example is shown in the figure below.

T(0)T(0) is the only subtree that is not nice. The function should return [0,1,1,1,1,1,1][0, 1, 1, 1, 1, 1, 1].

Constraints

  • 3N2000003 \le N \le 200\,000
  • 2M2000002 \le M \le 200\,000
  • 0P[i]<i0 \le P[i] \lt i (for all ii with 1i<N1 \le i \lt N)
  • 1C[i]M1 \le C[i] \le M (for all ii with 1i<N1 \le i \lt N)
  • P[0]=1P[0] = -1 and C[0]=0C[0] = 0

Subtasks

  1. (9 points) N8N \le 8 and M500M \le 500.
  2. (5 points) Edge ii connects node ii to node i1i-1. That is, for all ii with 1i<N1 \le i \lt N, we have P[i]=i1P[i] = i-1.
  3. (9 points) Except for node 00, every other node is connected either to node 00 or to some node that is connected to node 00.
    That is, for all ii with 1i<N1 \le i \lt N, either P[i]=0P[i] = 0 or P[P[i]]=0P[P[i]] = 0.
  4. (8 points) For each cc with 1cM1 \le c \le M, at most two edges have color cc.
  5. (14 points) N200N \le 200 and M500M \le 500.
  6. (14 points) N2000N \le 2\,000 and M=2M = 2.
  7. (12 points) N2000N \le 2\,000.
  8. (17 points) M=2M = 2.
  9. (12 points) No additional constraints.

Sample Grader

The sample grader reads input in the following format:

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

Let b[0],  b[1],  b[0], \; b[1], \; \ldots be the elements of the array returned by beechtree. The sample grader outputs your answer in the following format, on a single line:

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

Translated by ChatGPT 5