Description
There are $N$ students in SEATST. Each student represents exactly one country. After the contest, all the students obtained different scores.
Prabowo is about to post the leaderboard on the official website. For each student, the leaderboard lists their country, score, global rank, and country rank.
- The global rank of a student is defined as the number of students with a higher score than them.
- The country rank of a student is defined as the number of students within the same country with a higher score than them.
An example of a leaderboard is shown below:
| 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 |
Notice that, both the global ranks and country ranks start from $0$, and no ranks are skipped (either globally or within a country).
However, when the leaderboard was uploaded online, Prabowo forgot to post the students' countries and scores. For each student, we are given only their global rank and their rank within their country.
Prabowo is trying to salvage the situation and has tasked you to help them calculate two quantities:
- the number of pairs of distinct students that must belong to the same country, AND
- the number of pairs of distinct students that must belong to different countries.
Warning If there exist two assignments of students to countries that satisfy the global ranks and country ranks, where two students belong to the same country in one assignment and to two different countries in the other, then this pair of students must not be counted in either quantity.
Please help Prabowo!
Implementation Details
You should implement the following two procedures.
long long count_same_country(int N, std::vector<int> country_rank)
long long count_diff_country(int N, std::vector<int> country_rank)
- $N$: The number of students.
country_rank: An array of length $N$ representing the country ranks.country_rank[i]is the country rank of the student whose global rank is $i$, for all $0 \le i \le N - 1$.
The first procedure should return the number of unordered pairs of distinct students such that in all assignments of students to countries consistent with the leaderboard, the students in the pair belong to the same country.
The second procedure should return the number of unordered pairs of distinct students such that in all assignments of students to countries consistent with the leaderboard, the students in the pair do not belong to the same country.
Both procedures will be called at most once in each test case.
Constraints
- $1 \le N \le 1\;000\;000$.
- There exists an assignment of students to countries that satisfies
country_rank.
Subtasks
For the first 6 subtasks, only count_same_country will be called.
- ($3$ points) $N \leq 8$.
- ($6$ points)
country_rankcontains $0$ at most twice. - ($6$ points)
country_rankdoes not contain $2$. - ($3$ points) $N \leq 300$.
- ($3$ points) $N \leq 2000$.
- ($9$ points) No additional constraints.
For the last 6 subtasks, only count_diff_country will be called.
- ($7$ points) $N \leq 8$.
- ($14$ points)
country_rankcontains $0$ at most twice. - ($14$ points)
country_rankdoes not contain $2$. - ($7$ points) $N \leq 300$.
- ($7$ points) $N \leq 2000$.
- ($21$ points) No additional constraints.
Example
Consider the following procedure calls.
count_same_country(9, [0, 0, 1, 0, 2, 1, 3, 2, 1])
Let's assume that students $0$, $1$ and $3$ (here we are indexing students by their global rank for convenience) represent countries Singapore, Malaysia and Indonesia, respectively.
Then the following table lists all possible ways that this ranking could have been produced:
| Global Rank | Country Rank | Way 1 | Way 2 | Way 3 | Way 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 procedure 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 procedure should return $17$.
count_same_country(5, [0, 1, 0, 1, 2])
There are $2$ pairs of students here that must always belong in the same country: $(0,1)$ and $(2,3)$. Therefore, this procedure should return $2$.
count_diff_country(5, [0, 1, 0, 1, 2])
There are $4$ pairs of students who must belong to two different countries: $(0, 2), (0, 3), (1, 2), (1, 3)$. Therefore, this procedure should return $4$.
Sample Grader
Input Format:
N X
where X is either the string same or diff such that the function call being made is count_X_country and C[i] represents country_rank[i], for all $0 \leq i \leq N - 1$.
Output Format:
A single integer representing the output of count_X_country.