Completely Searching for Inversions (Codeforces 1866C)
Tags: DP, dfs and similar

Pak Chanek has a directed acyclic graph (a directed graph that does not have any cycles) containing $$$N$$$ vertices. Vertex $$$i$$$ has $$$S_i$$$ edges directed away from that vertex. The $$$j$$$-th edge of vertex $$$i$$$ that is directed away from it, is directed towards vertex $$$L_{i,j}$$$ and has an integer $$$W_{i,j}$$$ ($$$0\leq W_{i,j}\leq1$$$). Another information about the graph is that the graph is shaped in such a way such that each vertex can be reached from vertex $$$1$$$ via zero or more directed edges.Pak Chanek has an array $$$Z$$$ that is initially empty.Pak Chanek defines the function dfs as follows:// dfs from vertex ivoid dfs(int i) {    // iterate each edge of vertex i that is directed away from it    for(int j = 1; j <= S[i]; j++) {        Z.push_back(W[i][j]); // add the integer in the edge to the end of Z        dfs(L[i][j]); // recurse to the next vertex    }}Note that the function does not keep track of which vertices have been visited, so each vertex can be processed more than once.Let's say Pak Chanek does dfs(1) once. After that, Pak Chanek will get an array $$$Z$$$ containing some elements $$$0$$$ or $$$1$$$. Define an inversion in array $$$Z$$$ as a pair of indices $$$(x, y)$$$ ($$$x < y$$$) such that $$$Z_x > Z_y$$$. How many different inversions in $$$Z$$$ are there if Pak Chanek does dfs(1) once? Since the answer can be very big, output the answer modulo $$$998\,244\,353$$$.

Input: The first line contains a single integer $$$N$$$ ($$$2 \leq N \leq 10^5$$$) — the number of vertices in the graph. The following lines contain the description of each vertex from vertex $$$1$$$ to vertex $$$N$$$.The first line of each vertex $$$i$$$ contains a single integer $$$S_i$$$ ($$$0 \leq S_i \leq N-1$$$) — the number of edges directed away from vertex $$$i$$$.The $$$j$$$-th of the next $$$S_i$$$ lines of each vertex $$$i$$$ contains two integers $$$L_{i,j}$$$ and $$$W_{i,j}$$$ ($$$1 \leq L_{i,j} \leq N$$$; $$$0 \leq W_{i,j} \leq 1$$$) — an edge directed away from vertex $$$i$$$ that is directed towards vertex $$$L_{i,j}$$$ and has an integer $$$W_{i,j}$$$. For each $$$i$$$, the values of $$$L_{i,1}$$$, $$$L_{i,2}$$$, ..., $$$L_{i,S_i}$$$ are pairwise distinct. It is guaranteed that the sum of $$$S_i$$$ over all vertices does not exceed $$$2 \cdot 10^5$$$. There are no cycles in the graph. Each vertex can be reached from vertex $$$1$$$ via zero or more directed edges.

Output: An integer representing the number of inversions in $$$Z$$$ if Pak Chanek does dfs(1) once. Since the answer can be very big, output the answer modulo $$$998\,244\,353$$$.

Note: The following is the dfs(1) process on the graph. In the end, $$$Z=[0,1,0,1,1,0]$$$. All of its inversions are $$$(2,3)$$$, $$$(2,6)$$$, $$$(4,6)$$$, and $$$(5,6)$$$.

Sample Inputs:
5
2
4 0
3 1
0
1
2 0
2
3 1
5 1
0

Sample Outputs:
4

Solution (Python):
n = int(input())
edge = [[] for _ in range(n)]
for i in range(n):
    for _ in range(int(input())):
        (L, W) = map(int, input().split())
        edge[i].append((L - 1, W))
MOD = 998244353
dp = dict()
q = [[0] * 6]
while q:
    arr = q.pop()
    (i, j) = (arr[3], arr[4])
    if j - 1 >= 0 and (not arr[-1]):
        (zero, one, plus, _, _, _) = dp[edge[i][j - 1][0]]
        arr[2] = (arr[2] + arr[1] * zero + plus) % MOD
        arr[0] = (arr[0] + zero) % MOD
        arr[1] = (arr[1] + one) % MOD
    if j == len(edge[i]):
        dp[i] = tuple(arr)
        continue
    (nxt, w) = edge[i][j]
    arr[w] += 1
    if w == 0:
        arr[2] = (arr[2] + arr[1]) % MOD
    if nxt in dp:
        (zero, one, plus, _, _, _) = dp[nxt]
        arr[2] = (arr[2] + arr[1] * zero + plus) % MOD
        arr[0] = (arr[0] + zero) % MOD
        arr[1] = (arr[1] + one) % MOD
        arr[4] += 1
        arr[5] = 1
        q.append(arr)
    else:
        arr[4] += 1
        arr[5] = 0
        q.append(arr)
        q.append([0, 0, 0, nxt, 0, 0])
print(dp[0][2])

Solution (Python):
n = int(input())
edge = [[] for _ in range(n)]
for i in range(n):
    for _ in range(int(input())):
        (L, W) = map(int, input().split())
        edge[i].append((L - 1, W))
MOD = 998244353
dp = dict()
q = [[0] * 6]
while q:
    arr = q.pop()
    (i, j) = (arr[3], arr[4])
    if j - 1 >= 0 and (not arr[-1]):
        (zero, one, plus, _, _, _) = dp[edge[i][j - 1][0]]
        arr[2] = (arr[2] + arr[1] * zero + plus) % MOD
        arr[0] = (arr[0] + zero) % MOD
        arr[1] = (arr[1] + one) % MOD
    if j == len(edge[i]):
        dp[i] = tuple(arr)
        continue
    (nxt, w) = edge[i][j]
    arr[w] += 1
    if w == 0:
        arr[2] = (arr[2] + arr[1]) % MOD
    if nxt in dp:
        (zero, one, plus, _, _, _) = dp[nxt]
        arr[2] = (arr[2] + arr[1] * zero + plus) % MOD
        arr[0] = (arr[0] + zero) % MOD
        arr[1] = (arr[1] + one) % MOD
        arr[4] += 1
        arr[5] = 1
        q.append(arr)
    else:
        arr[4] += 1
        arr[5] = 0
        q.append(arr)
        q.append([0, 0, 0, nxt, 0, 0])
print(dp[0][2])

Solution (Python):
import sys
input = sys.stdin.readline

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

def readInt():
    return int(input())

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

def readStr():
    return input().strip()

def ops(a, b):
    a[-1] += a[1] * b[0] + b[-1]
    a[0] += b[0]
    a[1] += b[1]
    a[0] %= mo
    a[1] %= mo
    a[2] %= mo
    return a

def solve():
    n = readInt()
    graph = [[] for _ in range(n)]
    degree = [0] * n
    for i in range(n):
        for _ in range(readInt()):
            (j, w) = readInts()
            graph[i].append([w, j - 1, 1])
            graph[j - 1].append([w, i, 0])
            degree[i] += 1
    dp = [[0, 0, 0] for _ in range(n)]
    st = [i for i in range(n) if degree[i] == 0]
    while st:
        node = st.pop()
        curr = [0, 0, 0]
        for (w, j, flag) in graph[node]:
            if flag:
                curr = ops(curr, [0, 1, 0]) if w else ops(curr, [1, 0, 0])
                curr = ops(curr, dp[j])
        dp[node] = curr
        for (w, j, flag) in graph[node]:
            if not flag:
                degree[j] -= 1
                if degree[j] == 0:
                    st.append(j)
    return dp[0][-1]
mo = 998244353
print(solve())