Clearly, the kingdom can be modelled as a tree.

Subtask 1

If there is only one unique ~a_i~, then there must be one tribe, but this tribe could be at any node. Then, the answer is ~N~.

Otherwise, there must be two tribes. If either of the subgraphs induced by the nodes occupied by a particular tribe are disconnected, then the answer is ~0~. If not, then let ~s~ and ~t~ be the unique pair of two nodes occupied by tribes ~1~ and ~2~ respectively, such that ~s~ and ~t~ are adjacent. The answer is the number of pairs of nodes ~x, y~ such that ~x~ is occupied by tribe ~1~, ~Y~ is occupied by tribe ~2~, and ~|dist(s, x) - dist(y, t)| \le 1~. This can be computed by calculating the number of nodes occupied by tribe ~1~ at each distance from ~x~, and the number of nodes occupied by tribe ~2~ at each distance from ~y~.

Time complexity: ~\mathcal{O}(N)~

Subtask 2

Define a region as a maximal connected subgraph occupied by the same tribe. If there are multiple regions by the same tribe, the answer is ~0~. Our goal is to pick a headquarter node from each region such that for any two adjacent regions, the distances of headquarter nodes to the boundary of the regions differs by at most ~1~.

For nodes ~x~ and ~y~ in adjacent regions, with ~s~ being the node in ~x~'s region closest to ~y~, and ~t~ being the node in ~y~'s region closest to ~x~, define ~ok(x, y) = |dist(s, x) - dist(y, t)| \le 1~. Then we need ~ok(x, y)~ to hold for all chosen headquarters ~x, y~ in adjacent regions.

To avoid confusion, we will refer to the tree given in the input as the kingdom tree. We can generate a new tree, the region tree, by compressing regions in the kingdom tree into a single node. Root this region tree arbitrarily; let the root be ~a~. Denote ~c(r)~ as the set of children of region ~r~ in the region tree.

Let ~dp(x)~ be the number of kingdom states for the subtree of ~region(x)~ in the region tree, where ~x~ is chosen as the headquarters for its region. Then

~dp(x) = \prod\limits_{r \in c(region(x))} (\sum\limits_{y \in r} ([ok(x, y)] dp(y)))~

The answer is then ~\sum\limits_{region(x) = a} dp(x)~.

Time complexity: ~\mathcal{O}(N^2)~

Subtask 3

We need to speed up this DP. Consider using the dp values of one region to update the dp values for its parent in the region tree. We need a data structure over a tree which efficiently supports the following updates:

• Given a node ~x~, a distance ~d~ and a value ~v~, multiply the values for each node at a distance of ~d~ from ~x~ by the value ~v~.

This can be done efficiently by building a centroid tree over the parent region in the kingdom tree. We rely on the property that for any two nodes ~u~ and ~v~, the LCA of ~u~ and ~v~ in the centroid tree lies on the path between ~u~ and ~v~ in the kingdom tree. Thus, the set of nodes at distance ~d~ from ~x~ exactly corresponds to the set of nodes ~\{y | dist(x, LCA(x, y)) + dist(y, LCA(x, y)) = d\}~, where ~dist~ is the distance between nodes in the kingdom tree, and LCA is defined in the centroid tree. To handle an update, we can therefore iterate over all ancestors ~a~ of ~x~ in the centroid tree (possibly ~a = x~), and make a note to update all descendants ~y~ of ~a~ in the centroid tree such that ~dist(y, a) = d - dist(x, a)~ (note: we additionally require that ~a = LCA(x, y)~, but we will see how to handle this later). After all updates, for a given node ~f~, we can loop over all ancestors ~g~ of ~f~ in the centroid tree, and multiply all of the notes for node ~g~ tagged with distance ~dist(f, g)~.

There are multiple things to be careful about:

• If ~a \ne x~, let ~c~ be the child of ~a~ in the centroid tree which is closest to ~x~. In order for ~a = LCA(x, y)~ to hold, we require ~c~ not to lie on the path from ~a~ to ~y~. In other words, we only want to update descendants ~y~ of ~a~ in the centroid tree such that ~dist(y, a) = d - dist(x, a)~ and ~y~ does not lie in ~c~'s subtree in the centroid tree.
• Because of the above case, depending on implementation, it may be easier to separately consider paths which only go up the centroid tree, paths which only go down the centroid tree, and paths which go up and then down the centroid tree.
• For each node in a region, we need to make sure its DP value is the product of aggregated values from all of the children in the region tree. To clarify, consider an example where a parent region has a large depth, and one of its child regions has a small depth. Then some of the nodes in the parent region far away from the child region will be too far away to be updated by DP values in the child region; however, we actually want to set the DP values of these nodes to ~0~. To fix this, in addition to storing the number of kingdom states for each DP value, we can also store the number of values whose product contributed to the number of kingdom states. If this number of values does not match the number of child regions, then we can set the DP value to ~0~.
• You need to find an approach which only performs multiplications, since we cannot do modular division. There are multiple such approaches, such as doing sweep(s), or using prefix and suffix product arrays, but exact details are left as an exercise to the reader.

Time complexity: ~\mathcal{O}(N \log N)~

Bonus: there is an alternative solution which utilizes the linear merge trick to efficiently combine DP values of children, yielding an ~\mathcal{O}(N)~ solution.