Problem Description

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:

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.

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. 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.

Input Format

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 Format

Output one integer, the return value of count_X_country.

Hint

Samples

Consider the following function call:

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:

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