Counting Key‐Assignment Configurations in a Tree

You are given an undirected tree with n vertices numbered from 1 to n and an integer k (1 ≤ k ≤ n). Consider a sequence of k "key" vertices a₁, a₂, …, a_k (not necessarily distinct, though their roles are ordered) chosen from the vertices. For every vertex v in the tree, define its distance to a vertex u as the number of edges in the unique simple path connecting them. Now, for each vertex v let

\(f(v)=\min\{\,i: \operatorname{dis}(v,a_i)\ \text{is minimized}\}\)

where in the event of a tie (if two or more keys give the same distance) the smallest index is chosen. In other words, the "responsible key" for vertex v is the one (with smallest index in case of ties) among a₁, …, a_k minimizing the distance \(\operatorname{dis}(v,a_i)\). You are given the array of values f(1), f(2), …, f(n) and must count how many ordered sequences (a₁, a₂, …, a_k) can yield exactly the given function f when the rule above is applied. Since the answer may be large, output it modulo \(998244353\).

Important observations

• For \(1\le i\le k\), define the region \(S_i = \{v : f(v)=i\}\). In a valid configuration, each \(S_i\) must be nonempty and every \(S_i\) must induce a connected subgraph of the tree.
• The rule automatically forces a chain of border conditions. In any valid assignment the following must hold: for every edge \((u,v)\) in the tree, \(|f(u)-f(v)| \le 1\). Moreover, for every \(i\) with \(2\le i \le k\), if you look at the set of vertices in \(S_{i-1}\) having any neighbor in \(S_i\) (often called the border of \(S_{i-1}\) with respect to \(S_i\)), this set must be a singleton (i.e. it contains exactly one vertex). As it turns out, once the validity conditions are met the keys \(a_2, a_3, \dots, a_k\> are forced by these border conditions, and the only freedom is in choosing \(a_1\> arbitrarily from \(S_1\>.

Thus, one may show that if f is valid (i.e. for every edge the difference is at most 1, every region \(S_i\) is connected, and for each \(2\le i\le k\) the set \(\{v\in S_{i-1}: v\text{ is adjacent to some }u\in S_i\}\) is a singleton), then the number of key‐assignments producing f is exactly \(|S_1|\) modulo \(998244353\). Otherwise, there is no valid configuration and the answer is 0.

Input Format

1. The first line contains two integers \(n\) and \(k\) (\(1 \le k \le n \le 2 \cdot 10^5\)).
2. Each of the next \(n-1\) lines contains two integers \(u\) and \(v\) (\(1\le u,v\le n\)), denoting an edge of the tree.
3. The last line contains \(n\) integers: f(1), f(2), …, f(n), each between 1 and \(k\).

Output Format

Output a single integer: the number of ordered sequences \( (a_1, a_2, \dots, a_k) \) that produce the given function f by the rule described above, taken modulo \(998244353\).

inputFormat

The first line contains two integers n and k. The next n-1 lines each contain two integers u and v representing an edge of the tree. The last line contains n integers, the values of f(1), f(2), …, f(n).

outputFormat

A single integer — the number of valid key–assignment sequences modulo \(998244353\).

sample

3 1
1 2
2 3
1 1 1

3