Tree Elimination
Contest: Technocup 2020 - Elimination Round 4
Difficulty rating: 2900
Tags: dp, trees

Vasya has a tree with $$$n$$$ vertices numbered from $$$1$$$ to $$$n$$$, and $$$n - 1$$$ edges numbered from $$$1$$$ to $$$n - 1$$$. Initially each vertex contains a token with the number of the vertex written on it.

Vasya plays a game. He considers all edges of the tree by increasing of their indices. For every edge he acts as follows:

- If both endpoints of the edge contain a token, remove a token from one of the endpoints and write down its number.
- Otherwise, do nothing.

The result of the game is the sequence of numbers Vasya has written down. Note that there may be many possible resulting sequences depending on the choice of endpoints when tokens are removed.

Vasya has played for such a long time that he thinks he exhausted all possible resulting sequences he can obtain. He wants you to verify him by computing the number of distinct sequences modulo $$$998\,244\,353$$$.

Input format:
The first line contains a single integer $$$n$$$ ($$$2 \leq n \leq 2 \cdot 10^5$$$) — the number of vertices of the tree.

The next $$$n - 1$$$ lines describe edges of the tree. The $$$i$$$-th of these lines contains two integers $$$u_i, v_i$$$ ($$$1 \leq u_i, v_i \leq n$$$) — endpoints of the edge with index $$$i$$$. It is guaranteed that the given graph is indeed a tree.

Output format:
Print a single integer — the number of distinct sequences modulo $$$998\,244\,353$$$.

Constraints:
Time limit: 2.0 s
Memory limit: 512.0 MB

Note:
In the first sample case the distinct sequences are $$$(1), (2, 1), (2, 3, 1), (2, 3, 4, 1), (2, 3, 4, 5)$$$.

Int the second sample case the distinct sequences are $$$(2, 6, 5, 3), (2, 6, 5, 7), (2, 6, 7, 2), (2, 6, 7, 5), (2, 7, 3), (2, 7, 5), (7, 1, 3), (7, 1, 5), (7, 2, 3), (7, 2, 5)$$$.

Example input:
5
1 2
1 3
1 4
1 5

Example output:
5

Example input:
7
7 2
7 6
1 2
7 5
4 7
3 5

Example output:
10

Editorial:
First of all, counting different sequences is the same as counting the number of different playbacks of the elimination process (that is, different combinations of which token was removed on each step). Indeed, if we consider any resulting sequence, we can simulate the process and unambiguously determine which endpoint got its token removed on any step, skipping steps when no tokens can be removed.
To count different playbacks, we will use subtree dynamic programming. Let us consider a vertex $$$v$$$, and forget about all edges not incident to any vertex that is not in a subtree of $$$v$$$; that is, we are only considering edges in the subtree of $$$v$$$, as well as the edge between $$$v$$$ and its parent $$$p$$$ (for convenience, assume that the root vertex has an edge with index $$$n$$$ to a "virtual" parent). Note that we assume that $$$p$$$ can not be eliminated by means other than considering its edge to $$$v$$$.
As a shorthand, we will say "$$$v$$$ was compared with $$$u$$$" to mean "when the edge $$$(v, u)$$$ was considered, both its endpoints had their token", and "$$$v$$$ was killed by $$$u$$$" to mean "$$$v$$$ was compared with $$$u$$$ and lost its token on that step".
We will distinguish three types of playbacks in $$$v$$$'s subtree:
• $$$v$$$ was killed before comparing to $$$p$$$ (situation $$$0$$$);
• $$$v$$$ was killed by $$$p$$$ (situation $$$1$$$);
• $$$v$$$ killed $$$p$$$ (situation $$$2$$$).
We will write $$$dp_{v, s}$$$ the number of playbacks in the subtree of $$$v$$$ that correspond to the situation $$$s$$$.
Let $$$u_1, \ldots, u_k$$$ be the list of children of $$$v$$$ ordered by increasing of $$$\mathrm{index}(u_i, v)$$$, and let $$$d$$$ be the largest index such that $$$\mathrm{index}(u_d, v) < \mathrm{index}(v, p)$$$. If $$$\mathrm{index}(u_1, v) > \mathrm{index}(v, p)$$$, put $$$d = 0$$$.
Let us work out the recurrence relations for $$$dp_{v, s}$$$. For example, for $$$s = 0$$$ we must have that $$$v$$$ was killed by one of its children $$$u_i$$$ with $$$i \leq d$$$. For a fixed $$$i$$$, the playback should have proceeded as follows:
• all children $$$u_1, \ldots, u_i$$$ were either killed before comparing to $$$v$$$, or killed by $$$v$$$ (but they could not have survived comparing with $$$v$$$).
• $$$v$$$ was killed by $$$u_i$$$;
• all children $$$u_{i + 1}, \ldots, u_k$$$ were either killed before comparing to $$$v$$$, or "survived" the non-existent comparison to $$$v$$$ (but they could not have been killed by $$$v$$$).
Consequently, we have the formula $$$$$$dp_{v, 0} = \sum_{i = 1}^d \left(\prod_{j = 1}^{i = 1}(dp_{u_j, 0} + dp_{u_j, 1}) \times dp_{u_i, 1} \times \prod_{j = i + 1}^k (dp_{u_j, 0} + dp_{u_j, 2})\right).$$$$$$
Arguing in a similar way, we can further obtain
$$$$$$dp_{v, 1} = \prod_{j = 1}^{i = d}(dp_{u_j, 0} + dp_{u_j, 1}) \times \prod_{j = d + 1}^k (dp_{u_j, 0} + dp_{u_j, 2}),$$$$$$
$$$$$$dp_{v, 2} = \sum_{i = d + 1}^k \left(\prod_{j = 1}^{i = 1}(dp_{u_j, 0} + dp_{u_j, 1}) \times dp_{u_i, 1} \times \prod_{j = i + 1}^k (dp_{u_j, 0} + dp_{u_j, 2})\right) + \prod_{j = 1}^k (dp_{u_j, 0} + dp_{u_j, 1}).$$$$$$
In all these formulas we naturally assume that empty products are equal to $$$1$$$.
To compute these formulas fast enough we can use prefix products of $$$dp_{j, 0} + dp_{j, 1}$$$ and suffix products of $$$dp_{j, 0} + dp_{j, 2}$$$. Finally, the answer is equal to $$$dp_{root, 0} + dp_{root, 1}$$$ (either the root was killed, or it wasn't and we assume that it was killed by its virtual parent).
This solution can implemented in $$$O(n)$$$ since the edges are already given by increasing of their indices, but $$$O(n \log n)$$$ should also be enough.