P10437 [JOIST 2024] JOI Tour / JOI Tour

When submitting, please do not include `joitour.h`. Please do not submit using C++14 (GCC 9).

In the country of IOI, there are $N$ cities, numbered from $0$ to $N-1$, and $N-1$ roads, numbered from $0$ to $N-2$. Road $j\ (0\le j\le N-2)$ connects city $U_j$ and city $V_j$ bidirectionally. You can travel from one city to any other city via roads. Each city in the country of IOI has a restaurant. The restaurant type in city $i\ (0\le i\le N-1)$ is denoted by $F_i$. Specifically: - $F_i=0$: juice shop. - $F_i=1$: Japanese omelette shop. - $F_i=2$: ice cream shop. Ms. Rie is a tour guide in the country of IOI. She is planning a sightseeing program called the **JOI Tour**. In the JOI Tour, the plan is to visit the following $3$ types of restaurants in this order: 1. Choose a city $i_0\ (0\le i_0\le N-1)$ that has a juice shop, and start the trip from city $i_0$. 2. Visit the juice shop in city $i_0$. 3. Choose a city $i_1\ (0\le i_1\le N-1)$ that has a Japanese omelette shop, and take a bus from city $i_0$ to city $i_1$ along a shortest path. 4. Visit the Japanese omelette shop in city $i_1$. 5. Choose a city $i_2\ (0\le i_2\le N-1)$ that has an ice cream shop, and take a bus from city $i_1$ to city $i_2$ along a shortest path. 6. Visit the ice cream shop in city $i_2$. 7. End the trip in city $i_2$. To avoid boring the tourists, Ms. Rie decides to choose three cities $i_0,i_1,i_2$ such that they do not pass through the same road twice. We call such a JOI Tour a good one. To help her find an ideal plan, you are asked to compute the number of good JOI Tours. In other words, you need to find the number of triples $(i_0,i_1,i_2)$ that satisfy all of the following conditions: - The restaurant in city $i_0$ is a juice shop. - The restaurant in city $i_1$ is a Japanese omelette shop. - The restaurant in city $i_2$ is an ice cream shop. - While moving along shortest paths from city $i_0$ to $i_1$, and then from $i_1$ to $i_2$, you do not pass through the same road twice. In the country of IOI, there will be $Q$ events where restaurant types change. When the $(k+1)\ (0\le k\le Q-1)$-th event happens, two integers $X_k$ and $Y_k$ are given, satisfying $0\le X_k\le N-1$ and $0\le Y_k\le 2$. Then, the restaurant type in city $X_k$ becomes $Y_k$. That is, when $Y_k=0,1,2$, the restaurant type becomes a juice shop, a Japanese omelette shop, and an ice cream shop, respectively. After each event, you should immediately compute the latest number of good JOI Tours and tell Ms. Rie. Given the roads and restaurant information, compute the number of good JOI Tours after each type-change event. ### Implementation Details You need to implement the following functions. - `void init(int N, std::vector F, std::vector U, std::vector V, int Q)` - Use this function to provide the road and restaurant information. - This function is called only once at the beginning of the program. - The parameter `N` is the number of cities $N$. - The parameter `F` is an array of length $N$. `F[i]` ($0\le i\le N-1$) represents the restaurant type in city $i$, i.e. $F_i$. - The parameters `U` and `V` are arrays of length $N-1$. `U[j]` and `V[j]` ($0\le j\le N-2$) are the two cities $U_j$ and $V_j$ connected by road $j$. - The parameter `Q` is the number of restaurant type-change events $Q$. - `void change(int X, int Y)` - Use this function to provide a restaurant type-change event. - This function is called $Q$ times. - In the $(k+1)\ (0\le k\le Q-1)$-th call, the parameter `X` is the city index where the change happens, i.e. $X_k$. - In the $(k+1)\ (0\le k\le Q-1)$-th call, the parameter `Y` is the new restaurant type, i.e. $Y_k$. It is guaranteed that the new type is different from the old type. - `long long num_tours()` - This function is called in the following situations, for a total of $Q+1$ times. - After executing the function `init`. - After executing the function `change`. - This function should return the latest number of good JOI Tours.

**The following is the input format of the provided grader. You should not read any information from standard input.** The first line contains an integer $N$. The second line contains $N$ integers $F_0,\ldots,F_{N-1}$. The next $N-1$ lines each contain two integers $U_j,V_j$. The next line contains an integer $Q$. The next $Q$ lines each contain two integers $X_k,Y_k$.

**The following is the output format of the provided grader. You should not print any information to standard output.** After each call to the function `num_tours`, the provided grader outputs one line containing an integer, which is the return value of this function.

#### Sample Explanation 1 | Function Call | Return Value | | :-------------------------------------: | :----: | | `init(3, [0, 1, 2], [0, 1], [1, 2], 0)` | | | `num_tours()` | $1$ | There is exactly one good JOI Tour, which is $(i_0,i_1,i_2)=(0,1,2)$. Below is an explanation of why it satisfies the conditions of a good JOI Tour. - $F_0=0$, so the restaurant in city $0$ is a juice shop. - $F_1=1$, so the restaurant in city $1$ is a Japanese omelette shop. - $F_2=2$, so the restaurant in city $2$ is an ice cream shop. - When we move along shortest paths from city $0$ to $1$, and then from $1$ to $2$, we do not pass through the same road twice. Therefore, the return value of the first call to `num_tours()` is $1$. This sample satisfies the constraints of subtasks $1,2,3,4,6,7$. #### Sample Explanation 2 | Function Call | Return Value | | :-------------------------------------: | :----: | | `init(3, [0, 1, 2], [0, 1], [1, 2], 2)` | | | `num_tours()` | $1$ | | `change(2, 0)` | | | `num_tours()` | $0$ | | `change(0, 2)` | | | `num_tours()` | $1$ | Initially, there is one good JOI Tour, which is $(i_0,i_1,i_2)=(0,1,2)$. Therefore, the return value of the first call to `num_tours()` is $1$. In the first event, the restaurant in city $2$ changes from an ice cream shop to a juice shop. After this change, ice cream shops disappear from the country of IOI, so there are no good JOI Tours anymore. Therefore, the return value of the second call to `num_tours()` is $0$. In the second event, the restaurant in city $0$ changes from a juice shop to an ice cream shop. After this change, there is one good JOI Tour, which is $(i_0,i_1,i_2)=(2,1,0)$. Therefore, the return value of the third call to `num_tours()` is $1$. This sample satisfies the constraints of subtasks $1,2,4,6,7$. #### Sample Explanation 3 This sample satisfies the constraints of subtasks $1,2,5,6,7$. ### Important Notes - Your program may implement other functions for internal use, or use global variables. - Your program is not allowed to use standard input/output. Your program is not allowed to interact with other files in any way. However, your program may output debugging information to standard error output. ### Compilation and Test Run You can download the sample interactor in “File - Additional Files” to test your program. The additional files also include a sample program. The sample interactor is the file `grader.cpp`. To test your program, put `grader.cpp`, `joitour.cpp`, and `joitour.h` in the same directory, and compile with the following command. You may also run `compile.sh` to compile your program. ```bash g++ -std=gnu++20 -O2 -o grader grader.cpp joitour.cpp ``` If compilation succeeds, an executable file `grader` will be generated. Note that the interactor used for judging is different from the sample interactor. The sample interactor runs as a single process: it reads input from standard input and outputs results to standard output. ### Constraints - $3\le N\le 2\times 10^5$. - $0\le F_i\le 2\ (0\le i\le N-1)$. - $0\le U_j