Leaf Partition (Codeforces 1146F)
Tags: DP, trees

You are given a rooted tree with $$$n$$$ nodes, labeled from $$$1$$$ to $$$n$$$. The tree is rooted at node $$$1$$$. The parent of the $$$i$$$-th node is $$$p_i$$$. A leaf is node with no children. For a given set of leaves $$$L$$$, let $$$f(L)$$$ denote the smallest connected subgraph that contains all leaves $$$L$$$.You would like to partition the leaves such that for any two different sets $$$x, y$$$ of the partition, $$$f(x)$$$ and $$$f(y)$$$ are disjoint. Count the number of ways to partition the leaves, modulo $$$998244353$$$. Two ways are different if there are two leaves such that they are in the same set in one way but in different sets in the other.

Input: The first line contains an integer $$$n$$$ ($$$2 \leq n \leq 200\,000$$$) — the number of nodes in the tree.The next line contains $$$n-1$$$ integers $$$p_2, p_3, \ldots, p_n$$$ ($$$1 \leq p_i < i$$$). 

Output: Print a single integer, the number of ways to partition the leaves, modulo $$$998244353$$$.

Note: In the first example, the leaf nodes are $$$2,3,4,5$$$. The ways to partition the leaves are in the following image In the second example, the only leaf is node $$$10$$$ so there is only one partition. Note that node $$$1$$$ is not a leaf.

Sample Inputs:
5
1 1 1 1

10
1 2 3 4 5 6 7 8 9

Sample Outputs:
12

1

Solution (Python):
import sys
input = lambda : sys.stdin.readline().rstrip('\r\n')
sint = lambda : int(input())
mint = lambda : map(int, input().split())
aint = lambda : list(map(int, input().split()))
T = 1
mod = 998244353
for _ in range(T):
    n = sint()
    a = aint()
    g = [[] for _ in range(n + 1)]
    for i in range(n - 1):
        (u, v) = (i + 2, a[i])
        g[v].append(u)
    dp = [[0] * (n + 1) for _ in range(3)]
    for c in range(n, 0, -1):
        if len(g[c]) == 0:
            dp[2][c] = 1
            continue
        dp[0][c] = 1
        for i in g[c]:
            dp[2][c] = (dp[2][c] * (dp[0][i] + dp[1][i] + 2 * dp[2][i]) + dp[1][c] * (dp[1][i] + dp[2][i])) % mod
            dp[1][c] = (dp[1][c] * (dp[0][i] + dp[2][i]) + dp[0][c] * (dp[1][i] + dp[2][i])) % mod
            dp[0][c] = dp[0][c] * (dp[0][i] + dp[2][i]) % mod
    print((dp[0][1] + dp[2][1]) % mod)

Solution (Python):
(standard_input, packages, output_together) = (1, 1, 1)
(dfs, hashing, read_from_file) = (0, 0, 0)
de = 1
if 1:
    if standard_input:
        import io, os, sys
        input = lambda : sys.stdin.readline().strip()
        inf = float('inf')

        def I():
            return input()

        def II():
            return int(input())

        def MII():
            return map(int, input().split())

        def LI():
            return list(input().split())

        def LII():
            return list(map(int, input().split()))

        def LFI():
            return list(map(float, input().split()))

        def GMI():
            return map(lambda x: int(x) - 1, input().split())

        def LGMI():
            return list(map(lambda x: int(x) - 1, input().split()))
    if packages:
        from io import BytesIO, IOBase
        import math
        import random
        import os
        import bisect
        import typing
        from collections import Counter, defaultdict, deque
        from copy import deepcopy
        from functools import cmp_to_key, lru_cache, reduce
        from heapq import merge, heapify, heappop, heappush, heappushpop, nlargest, nsmallest
        from itertools import accumulate, combinations, permutations, count, product
        from operator import add, iand, ior, itemgetter, mul, xor
        from string import ascii_lowercase, ascii_uppercase, ascii_letters
        from typing import *
        BUFSIZE = 4096
    if output_together:

        class FastIO(IOBase):
            newlines = 0

            def __init__(self, file):
                self._fd = file.fileno()
                self.buffer = BytesIO()
                self.writable = 'x' in file.mode or 'r' not in file.mode
                self.write = self.buffer.write if self.writable else None

            def read(self):
                while True:
                    b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
                    if not b:
                        break
                    ptr = self.buffer.tell()
                    (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
                self.newlines = 0
                return self.buffer.read()

            def readline(self):
                while self.newlines == 0:
                    b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
                    self.newlines = b.count(b'\n') + (not b)
                    ptr = self.buffer.tell()
                    (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
                self.newlines -= 1
                return self.buffer.readline()

            def flush(self):
                if self.writable:
                    os.write(self._fd, self.buffer.getvalue())
                    (self.buffer.truncate(0), self.buffer.seek(0))

        class IOWrapper(IOBase):

            def __init__(self, file):
                self.buffer = FastIO(file)
                self.flush = self.buffer.flush
                self.writable = self.buffer.writable
                self.write = lambda s: self.buffer.write(s.encode('ascii'))
                self.read = lambda : self.buffer.read().decode('ascii')
                self.readline = lambda : self.buffer.readline().decode('ascii')
        sys.stdout = IOWrapper(sys.stdout)
    if dfs:
        from types import GeneratorType

        def bootstrap(f, stack=[]):

            def wrappedfunc(*args, **kwargs):
                if stack:
                    return f(*args, **kwargs)
                else:
                    to = f(*args, **kwargs)
                    while True:
                        if type(to) is GeneratorType:
                            stack.append(to)
                            to = next(to)
                        else:
                            stack.pop()
                            if not stack:
                                break
                            to = stack[-1].send(to)
                    return to
            return wrappedfunc
    if hashing:
        RANDOM = random.getrandbits(20)

        class Wrapper(int):

            def __init__(self, x):
                int.__init__(x)

            def __hash__(self):
                return super(Wrapper, self).__hash__() ^ RANDOM
    if read_from_file:
        file = open('input.txt', 'r').readline().strip()[1:-1]
        fin = open(file, 'r')
        input = lambda : fin.readline().strip()
        output_file = open('output.txt', 'w')

        def fprint(*args, **kwargs):
            print(*args, **kwargs, file=output_file)
    if de:

        def debug(*args, **kwargs):
            print('\x1b[92m', end='')
            print(*args, **kwargs)
            print('\x1b[0m', end='')

def main():
    n = II()
    parent = [-1] + LGMI()
    mod = 998244353
    dp = [[-1] * n for _ in range(3)]
    for i in range(n - 1, -1, -1):
        if dp[0][i] == -1:
            dp[0][i] = dp[1][i] = 0
            dp[2][i] = 1
        if i:
            p = parent[i]
            if dp[0][p] == -1:
                dp[0][p] = 1
                dp[1][p] = dp[2][p] = 0
            dp[2][p] *= dp[0][i] + dp[1][i] + 2 * dp[2][i]
            dp[2][p] += dp[1][p] * (dp[1][i] + dp[2][i])
            dp[2][p] %= mod
            dp[1][p] *= dp[0][i] + dp[2][i]
            dp[1][p] += dp[0][p] * (dp[1][i] + dp[2][i])
            dp[1][p] %= mod
            dp[0][p] *= dp[0][i] + dp[2][i]
            dp[0][p] %= mod
    print((dp[0][0] + dp[2][0]) % mod)
    return
t = 1
for _ in range(t):
    main()

Solution (Python):
import sys
input = lambda : sys.stdin.readline().rstrip('\r\n')
sint = lambda : int(input())
mint = lambda : map(int, input().split())
aint = lambda : list(map(int, input().split()))
T = 1
mod = 998244353
for _ in range(T):
    n = sint()
    a = aint()
    g = [[] for _ in range(n + 1)]
    for i in range(n - 1):
        (u, v) = (i + 2, a[i])
        g[u].append(v)
        g[v].append(u)
    par = [0] * (n + 1)
    dp = [[0] * (n + 1) for _ in range(3)]
    Q = [1]
    ans = 0
    while Q:
        c = Q.pop()
        if c > 0:
            Q.append(~c)
            for i in g[c]:
                if i == par[c]:
                    continue
                Q.append(i)
                par[i] = c
        else:
            c = ~c
            if c > 1 and len(g[c]) == 1:
                dp[2][c] = 1
                continue
            dp[0][c] = 1
            for i in g[c]:
                if i == par[c]:
                    continue
                dp[2][c] = (dp[2][c] * (dp[0][i] + dp[1][i] + 2 * dp[2][i]) + dp[1][c] * (dp[1][i] + dp[2][i])) % mod
                dp[1][c] = (dp[1][c] * (dp[0][i] + dp[2][i]) + dp[0][c] * (dp[1][i] + dp[2][i])) % mod
                dp[0][c] = dp[0][c] * (dp[0][i] + dp[2][i]) % mod
    print((dp[0][1] + dp[2][1]) % mod)