thảo luận Leetcode + Codeforces, Competitive programming contest. Đường tới Guardian + Candidate Master.

from collections import defaultdict
from typing import List

class Solution:
    def countValidSubsets(self, parent: List[int], nums: List[int], k: int) -> int:
        MOD = 10**9 + 7
        graph = defaultdict(list)
        for i, val in enumerate(parent):
            if val != -1:
                graph[val].append(i)

        def dfs(node, p):
            # Base state for the current node BEFORE looking at any neighbors
            # includeNode tracks if we pick 'node'. Starts with just its own value.
            includeNode = defaultdict(int)
            includeNode[nums[node] % k] = 1
            
            # excludeNode tracks if we don't pick 'node'. Starts with an empty subset (sum 0).
            excludeNode = defaultdict(int)
            excludeNode[0] = 1

            for neighbor in graph[node]:
                if neighbor == p:
                    continue

                includeNeighbor, excludeNeighbor = dfs(neighbor, node)
                
                # Temporary containers to hold the multiplied combinations safely
                nextInclude = defaultdict(int)
                nextExclude = defaultdict(int)

                # 1. If 'node' is INCLUDED, the neighbor MUST be EXCLUDED
                for k_parent, v_parent in includeNode.items():
                    for k_child, v_child in excludeNeighbor.items():
                        rem = (k_parent + k_child) % k
                        nextInclude[rem] = (nextInclude[rem] + v_parent * v_child) % MOD

                # 2. If 'node' is EXCLUDED, the neighbor can be INCLUDED OR EXCLUDED
                for k_parent, v_parent in excludeNode.items():
                    # Combined total ways for the child branch at remainder k_child
                    all_child_keys = set(includeNeighbor.keys()) | set(excludeNeighbor.keys())
                    
                    for k_child in all_child_keys:
                        v_child = (includeNeighbor[k_child] + excludeNeighbor[k_child]) % MOD
                        rem = (k_parent + k_child) % k
                        nextExclude[rem] = (nextExclude[rem] + v_parent * v_child) % MOD

                # Update our current node's profiles with the new multiplied states
                includeNode = nextInclude
                excludeNode = nextExclude

            return includeNode, excludeNode
                
        # Run DFS from the root
        rootInclude, rootExclude = dfs(0, -1)
        
        # Total valid subsets have a remainder of 0.
        # We subtract 1 to remove the "completely empty subset" (picking nothing at all)
        ans = (rootInclude[0] + rootExclude[0] - 1) % MOD
        return ans©leetcode