P16704 [SEATST 2026] National Rankings / Country Ranks

There are $N$ students participating in the SEATST contest. Each student represents a country. After the contest ends, all students receive **distinct** scores. Prabowo is going to publish the ranking table on the official website. For each student, the ranking table lists their country, score, global rank, and country rank. - A student's global rank is defined as the number of people whose score is higher than this student's score. - A student's **country rank** is defined as the number of people who are from the same country as this student and whose score is higher than this student's score. An example ranking table is as follows: | Country | Score | Global Rank | Country Rank | |:-:|:-:|:-:|:-:| | Singapore | $574$ | $0$ | $0$ | | Malaysia | $483$ | $1$ | $0$ | | Singapore | $466$ | $2$ | $1$ | | Indonesia | $460$ | $3$ | $0$ | | Singapore | $458$ | $4$ | $2$ | | Malaysia | $454$ | $5$ | $1$ | | Singapore | $448$ | $6$ | $3$ | | Malaysia | $440$ | $7$ | $2$ | | Indonesia | $438$ | $8$ | $1$ | Note that both global rank and country rank start from $0$, and the ranks never skip any numbers (for both global rank and country rank). However, when the ranking table was uploaded online, Prabowo forgot to publish the students' countries and scores. For each student, we only know their global rank and country rank. Prabowo tries to save the situation, and gives you a task to help him compute the following two quantities: - the number of pairs of students that must belong to the **same** country, and - the number of pairs of students that must belong to **different** countries. :::warning[Warning]{open} If there exist two assignments consistent with the global ranks and country ranks such that two students are in the same country in one assignment, but in different countries in the other assignment, then this pair of students should not be counted in either of the two quantities above. ::: Please help Prabowo. ### Implementation Details You need to implement the following two functions. ```cpp long long count_same_country(int N, std::vector country_rank) long long count_diff_country(int N, std::vector country_rank) ``` - $N$: the number of students. - `country_rank`: an array of length $N$ representing the country ranks. For all $0 \le i \le N - 1$, `country_rank[i]` is the country rank of the student whose global rank is $i$. The first function should return the number of unordered pairs of distinct students such that, in all country assignments consistent with the ranking table, the two students always belong to the same country. The second function should return the number of unordered pairs of distinct students such that, in all country assignments consistent with the ranking table, the two students always belong to different countries. In each testdata, each of these two functions will be called at most once.

The input format is: ``` N X C[0] C[1] ... C[N-1] ``` Here, `X` can be the string `same` or `diff`, forming a call to the function `count_X_country`. For all $0 \le i \le N - 1$, `C[i]` denotes `country_rank[i]`.

Output one integer, the return value of `count_X_country`.

### Samples Consider the following function call: ```cpp count_same_country(9, [0, 0, 1, 0, 2, 1, 3, 2, 1]) ``` Assume that students $0$, $1$, and $3$ (for convenience, here we number students by their global ranks) represent Singapore, Malaysia, and Indonesia respectively. Then, the table below lists all assignments that can produce these ranks: | Global Rank | Country Rank | Assignment 1 | Assignment 2 | Assignment 3 | Assignment 4 | |:-:|:-:|:-:|:-:|:-:|:-:| | $0$ | $0$ | Singapore | Singapore | Singapore | Singapore | | $1$ | $0$ | Malaysia | Malaysia | Malaysia | Malaysia | | $2$ | $1$ | Singapore | Singapore | Malaysia | Malaysia | | $3$ | $0$ | Indonesia | Indonesia | Indonesia | Indonesia | | $4$ | $2$ | Singapore | Singapore | Malaysia | Malaysia | | $5$ | $1$ | Malaysia | Indonesia | Singapore | Indonesia | | $6$ | $3$ | Singapore | Singapore | Malaysia | Malaysia | | $7$ | $2$ | Malaysia | Indonesia | Singapore | Indonesia | | $8$ | $1$ | Indonesia | Malaysia | Indonesia | Singapore | There are $4$ pairs of students that must always belong to the same country: $(2, 4)$, $(2, 6)$, $(4, 6)$, and $(5, 7)$. Therefore, this function should return $4$. `count_diff_country(9, [0, 0, 1, 0, 2, 1, 3, 2, 1])` There are $17$ pairs of students that must always belong to different countries: $(0, 1)$, $(0, 3)$, $(1, 3)$, $(2, 3)$, $(2, 5)$, $(2, 7)$, $(2, 8)$, $(3, 4)$, $(3, 6)$, $(4, 5)$, $(4, 7)$, $(4, 8)$, $(5, 6)$, $(5, 8)$, $(6, 7)$, $(6, 8)$, $(7, 8)$. Therefore, this function should return $17$. `count_same_country(5, [0, 1, 0, 1, 2])` Here there are $2$ pairs of students that must always belong to the same country: $(0, 1)$ and $(2, 3)$. Therefore, this function should return $2$. `count_diff_country(5, [0, 1, 0, 1, 2])` There are $4$ pairs of students that must belong to two different countries: $(0, 2)$, $(0, 3)$, $(1, 2)$, $(1, 3)$. Therefore, this function should return $4$. ### Constraints - $1 \le N \le 1\ 000\ 000$. - It is guaranteed that there exists at least one country assignment satisfying `country_rank`. ### Subtasks For the first $6$ subtasks, only `count_same_country` will be called. 1. ($3$ points) $N \le 8$. 2. ($6$ points) `country_rank` contains at most two $0$. 3. ($6$ points) `country_rank` does not contain $2$. 4. ($3$ points) $N \le 300$. 5. ($3$ points) $N \le 2000$. 6. ($9$ points) No additional constraints. For the last $6$ subtasks, only `count_diff_country` will be called. 7. ($7$ points) $N \le 8$. 8. ($14$ points) `country_rank` contains at most two $0$. 9. ($14$ points) `country_rank` does not contain $2$. 10. ($7$ points) $N \le 300$. 11. ($7$ points) $N \le 2000$. 12. ($21$ points) No additional constraints. Translated by ChatGPT 5