Yet Another Contest 9 P6 - Tribes

View as PDF

Submit solution


Points: 25 (partial)
Time limit: 5.0s
Memory limit: 1G

Author:
Problem types

On his holiday, Josh finds himself in a kingdom with N villages, which he labels from 1 to N. The kingdom is connected by N - 1 roads with equal length, with the i-th road connecting villages u_i and v_i, such that each village can be reached from every other village. The distance between two villages is the minimum number of roads needed to travel between them.

Josh's tour guide tells him that there are also some positive number of tribes, each denoted with a positive integer, and each having a headquarter in exactly one village. No two tribes have the same headquarter. Additionally, each village is occupied by exactly one tribe - the tribe whose headquarters are the shortest distance away from the village. If there are multiple tribes with headquarters being the shortest distance away from the village, the village could be occupied by any of those tribes.

After exploring each village, Josh determines that the i-th village is given occupied by tribe a_i, but he doesn't know where any of the headquarters are. He additionally notes that for any two villages occupied by the same tribe x, it is possible to travel between the two villages along the roads, only travelling via villages occupied by tribe x.

Josh would like to know the number of different possible states of the kingdom. Two kingdoms states are considered different if there exists a village satisfying one of the two conditions below:

  • Both kingdom states have a headquarter at that village, but the headquarters belong to different tribes.
  • One kingdom state has a headquarter at that village, and the other kingdom state doesn't.

Since this number might be very large, output it modulo M. Note that there might be no possible states, in which case the answer is 0.

Note: it is not required that the tribes present are labelled from 1 onwards. For example, it is possible that tribe 2 occupies the entire kingdom.

Constraints

1 \le N \le 5 \times 10 ^ 5

2 \le M \le 10^9

1 \le a_i \le 5 \times 10 ^ 5

1 \le u_i, v_i \le N

All villages can be reached from all other villages by travelling along the roads.

For any two villages occupied by the same tribe x, it is possible to travel between the two villages along the roads, only travelling via villages occupied by tribe x.

Subtask 1 [15%]

1 \le a_i \le 2

Subtask 2 [15%]

1 \le N \le 2000

Subtask 3 [70%]

No additional constraints.

Input Specification

The first line contains two space-separated integers, N and M.

The second line contains N space-separated integers, a_1, a_2, \dots, a_N.

The i-th of the following N - 1 lines contains two space-separated integers, u_i, v_i.

Output Specification

Output the number of possible kingdom states modulo M.

Sample Input

5 10
1 1 2 2 2
1 2
1 3
3 4
3 5

Sample Output

6

Explanation

For this kingdom, tribe 1 could have its headquarters in villages 1 or 2, and tribe 2 could have its headquarters in village 3, 4 or 5. Furthermore, all of these combinations yield a valid kingdom state. For example, if tribe 1 has its headquarters in village 2, and tribe 2 has its headquarters in village 3, then:

  • Village 1 is equidistant to both headquarters, so it could feasibly be occupied by tribe 1.
  • Village 2 is closest to the headquarters located in village 2, which belongs to tribe 1.
  • Villages 3, 4 and 5 are closest to the headquarters located in village 3, which belongs to tribe 2.

Comments

There are no comments at the moment.