Forgotten Land
Time limit3sMemory limit512 MB
Sum over all partitions of a tree's vertices into arbitrary sets of the total language difficulty, where a set's difficulty depends on the union of languages on its vertices and on paths between them.
- Level
Hard9 of 10
- Topics
- Dynamic programming, Tree, Combinatorics, Math
- Solved
- No attempts yet
Problem
Fytelandian scientists are famous for their love of historical research. Recently they discovered the remains of n cities that definitely belonged to a great past civilization.
All the discovered cities were connected by exactly n − 1 roads, and there was exactly one way to travel between any pair of cities using these roads. The scientists found many writings in each city and concluded that the people of this civilization spoke k different languages. In city v, people spoke language av.
The cities formed several alliances, and each city belonged to exactly one alliance. However, the exact composition of the alliances is unknown.
An alliance is an arbitrary set of cities c1, c2, . . . , cm, which may or may not be connected by roads. To rule an alliance, it had to be possible to connect different cities whose people speak different languages. The head of the alliance had to support every language spoken either in some city ci of the alliance or in some city on the shortest path between two cities ci and cj of the alliance.
The more languages supported in an alliance, the harder the translator's job was. A good translator had to speak all of them. Fortunately, many languages were similar to each other, and the more languages the translator already knew, the easier it was to learn a new one. Each new language took half as much time to learn as the previous one.
The language difficulty of an alliance is the time the translator needed to learn all languages supported in the alliance. The first language took 2k units of time to learn, so the language difficulty of the alliance equals 2k + 2k−1 + . . . + 2k+1−t, where t is the number of distinct languages supported in the alliance. Since there are k languages in total, this sum is always an integer.
An alliance partition is a partition of all cities into several alliances. Two alliance partitions are considered different if there exist two cities u and v that belong to the same alliance in one partition and to different alliances in the other.
The plausibility of an alliance partition is the sum of the language difficulties of all alliances in the partition.
You need to compute the sum of the plausibilities of all possible alliance partitions. Since this number can be very large, find only its remainder modulo 998 244 353.
Input
The first line of the input contains two integers n and k: the number of cities and the number of languages (1 ≤ n ≤ 5000; 1 ≤ k ≤ 10).
The second line contains n integers a1, a2, . . . , an: the languages spoken in the cities (1 ≤ ai ≤ k).
Each of the following n − 1 lines contains two integers ui and vi: the cities connected by the i-th road (1 ≤ ui, vi ≤ n).
Output
Output the sum of the plausibilities of all possible alliance partitions modulo 998 244 353.