1866C - Completely Searching for Inversions - CodeForces Solution

#include<iostream>
#include<map>
#include<vector>
#include<cstring>
#include<unordered_set>
#include<unordered_map>
using namespace std;
const static int N=1e5+5;
long long dp[N][3];
int mod=998244353;
vector<pair<int,int>> edge[N];
bool vis[N];
long long cnt1=0;
long long ans=0; 

void dfs(int u){
//     long long tmp1=0;
    for(auto &pp:edge[u]){
        if(!vis[pp.first])
            dfs(pp.first);
//         if(pp.second)tmp1++;
        dp[u][1] = (dp[u][1] + (pp.second))%mod;
        dp[u][2] = ((dp[u][2]+ dp[pp.first][2])%mod + (dp[pp.first][0] + (pp.second == 0))*dp[u][1]%mod)%mod;
        dp[u][0] = ((dp[u][0] +dp[pp.first][0])%mod + (pp.second==0))%mod;
        dp[u][1] = ((dp[u][1] +dp[pp.first][1])%mod)%mod;
        
//         tmp1=(dp[pp.first][1]+tmp1)%mod;
    }
    vis[u] = true;
}
int main(){
    int n;
    cin>>n;
    for(int i=1;i<=n;i++){
        int m;
        int a,b;
        cin>>m;
        for(int j=0;j<m;j++){
            cin>>a>>b;
            edge[i].push_back({a,b});
        }
    }
    dfs(1);
    cout<<dp[1][2];

    
}