P9603 [IOI 2023] Beech Tree

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.

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 $N$ **nodes** and $N-1$ **edges**. The nodes are numbered from $0$ to $N-1$, and the edges are numbered from $1$ to $N-1$. Each edge connects two different nodes in the tree. Specifically, edge $i$ ($1 \le i \lt N$) connects node $i$ to node $P[i]$, where $0 \le P[i] \lt i$. Node $P[i]$ is called the **parent** of node $i$, and node $i$ is called a **child** of node $P[i]$. Each edge has a color. There are $M$ possible colors, numbered from $1$ to $M$. The color of edge $i$ is $C[i]$. Different edges may have the same color. Note that in the definition above, the case $i = 0$ does not correspond to an edge in the tree. For convenience, we set $P[0] = -1$ and $C[0] = 0$. For example, suppose Ős Vezér has $N = 18$ nodes and $M = 3$ possible colors, and thus $17$ 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 $r$ with $0 \le r \lt N$, the subtree of node $r$ is a set of nodes $T(r)$ with the following properties: * Node $r$ is in $T(r)$. * If a node $x$ is in $T(r)$, then all children of $x$ are also in $T(r)$. * No other nodes are in $T(r)$. The size of the set $T(r)$ is denoted by $|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 $r$, and a permutation $v_0, v_1, \ldots, v_{|T(r)|-1}$ of the nodes in the subtree $T(r)$. For each $i$ with $1 \le i \lt |T(r)|$, let $f(i)$ be the number of occurrences of the color $C[v_i]$ in the color sequence of length $i-1$, namely $C[v_1], C[v_2], \ldots, C[v_{i-1}]$. (Note that $f(1)$ is always $0$, because the sequence in its definition is empty.) The permutation $v_0, v_1, \ldots, v_{|T(r)|-1}$ is called a **nice permutation** if and only if the following conditions hold: * $v_0 = r$. * For each $i$ with $1 \le i \lt |T(r)|$, the parent of node $v_i$ is $v_{f(i)}$. For each $r$ with $0 \le r \lt N$, the subtree $T(r)$ is called a **nice subtree** if and only if there exists some nice permutation of the nodes in $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)$ and $T(3)$ are not nice. Subtree $T(14)$ is nice because it contains only one node. Next, we will show that subtree $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)$. 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**. * $v_0 = 1$. * $f(1) = 0$, because $C[v_1] = C[4] = 1$ occurs $0$ times in the sequence $[\,]$. * Accordingly, the parent of $v_1$ is $v_0$. That is, the parent of $4$ is $1$. (Formally, $P[4] = 1$.) * $f(2) = 0$, because $C[v_2] = C[5] = 2$ occurs $0$ times in the sequence $[1]$. * Accordingly, the parent of $v_2$ is $v_0$. That is, the parent of $5$ is $1$. * $f(3) = 1$, because $C[v_3] = C[12] = 1$ occurs $1$ time in the sequence $[1, 2]$. * Accordingly, the parent of $v_3$ is $v_1$. That is, the parent of $12$ is $4$. * $f(4) = 1$, because $C[v_4] = C[13] = 2$ occurs $1$ time in the sequence $[1, 2, 1]$. * Accordingly, the parent of $v_4$ is $v_1$. That is, the parent of $13$ is $4$. * $f(5) = 0$, because $C[v_5] = C[6] = 3$ occurs $0$ times in the sequence $[1, 2, 1, 2]$. * Accordingly, the parent of $v_5$ is $v_0$. That is, the parent of $6$ is $1$. * $f(6) = 2$, because $C[v_6] = C[14] = 2$ occurs $2$ times in the sequence $[1, 2, 1, 2, 3]$. * Accordingly, the parent of $v_6$ is $v_2$. That is, the parent of $14$ is $5$. Since we can find a **nice permutation** for the nodes in $T(1)$, the subtree $T(1)$ is therefore a **nice subtree**. Your task is to help Árpád determine whether each subtree of Ős Vezér is nice.

You need to implement the following function. ``` int[] beechtree(int N, int M, int[] P, int[] C) ``` * $N$: the number of nodes in the tree. * $M$: the number of possible edge colors in the tree. * $P$, $C$: two arrays of length $N$ describing the edges of the tree. * The function should return an array $b$ of length $N$. For each $r$ with $0 \le r \lt N$, if $T(r)$ is nice, then $b[r]$ should be $1$, otherwise it should be $0$. * This function is called exactly once for each test case.

N/A

### 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(2)$, and $T(3)$ each contain a single node, so they are all nice. $T(0)$ is not nice. Therefore, the function should return $[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)$ is the only subtree that is not nice. The function should return $[0, 1, 1, 1, 1, 1, 1]$. ### Constraints * $3 \le N \le 200\,000$ * $2 \le M \le 200\,000$ * $0 \le P[i] \lt i$ (for all $i$ with $1 \le i \lt N$) * $1 \le C[i] \le M$ (for all $i$ with $1 \le i \lt N$) * $P[0] = -1$ and $C[0] = 0$ ### Subtasks 1. (9 points) $N \le 8$ and $M \le 500$. 1. (5 points) Edge $i$ connects node $i$ to node $i-1$. That is, for all $i$ with $1 \le i \lt N$, we have $P[i] = i-1$. 1. (9 points) Except for node $0$, every other node is connected either to node $0$ or to some node that is connected to node $0$. That is, for all $i$ with $1 \le i \lt N$, either $P[i] = 0$ or $P[P[i]] = 0$. 1. (8 points) For each $c$ with $1 \le c \le M$, at most two edges have color $c$. 1. (14 points) $N \le 200$ and $M \le 500$. 1. (14 points) $N \le 2\,000$ and $M = 2$. 1. (12 points) $N \le 2\,000$. 1. (17 points) $M = 2$. 1. (12 points) No additional constraints. ### Sample Grader The sample grader reads input in the following format: * Line $1$: $N \; M$ * Line $2$: $P[0] \; P[1] \; \ldots \; P[N-1]$ * Line $3$: $C[0] \; C[1] \; \ldots \; C[N-1]$ Let $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 $1$: $b[0] \; b[1] \; \ldots$ Translated by ChatGPT 5