LeetCode Problem 3939

Difficulty: 🔴 Hard
Topics: —

Solution

LeetCode 3939 - Count Non Adjacent Subsets in a Rooted Tree

Problem Understanding

We are given a rooted tree with n nodes. The tree is represented by the parent array, where parent[i] tells us the parent of node i. Node 0 is always the root because parent[0] = -1.

Each node also has a value stored in nums[i].

We want to count how many non-empty subsets of nodes satisfy two conditions simultaneously:

1. The sum of the selected node values is divisible by k.
2. No two selected nodes are directly connected by an edge. In other words, a node and its parent cannot both belong to the subset.

The second condition means the chosen nodes must form an independent set in the tree.

The answer can be extremely large, so we return it modulo 10^9 + 7.

The constraints are important:

• n ≤ 1000
• k ≤ 100
• nums[i] can be very large, up to 10^9

The relatively small value of k is the key observation. Since divisibility only depends on the remainder modulo k, we can store information using only k possible residue classes.

Some important edge cases include:

• A tree containing only one node.
• Cases where no valid subset exists.
• Cases where the only divisible subset is the empty subset, which must not be counted.
• Long chains, where adjacency restrictions are strongest.
• Star-shaped trees, where many children can be chosen simultaneously.
• Large node values, requiring us to reduce values modulo k.

Approaches

Brute Force

A brute force solution would examine every possible subset of nodes.

For each subset, we would:

1. Check whether any selected node and its parent are both included.
2. Compute the sum of selected values.
3. Verify divisibility by k.

Since there are 2^n subsets, this approach requires exponential time.

Even for n = 50, this would already be infeasible, and the actual constraint is n = 1000.

Therefore, brute force is impossible.

Key Insight

The problem combines two properties:

• The adjacency restriction is local to tree edges.
• Divisibility only depends on the sum modulo k.

This naturally suggests a tree dynamic programming solution.

For each node, we want to know:

• How many valid configurations exist inside its subtree.
• What residue modulo k each configuration produces.
• Whether the current node itself is selected.

The selection state is necessary because selecting a node restricts what its children are allowed to do.

This leads to a classic tree DP for independent sets, augmented with residue tracking modulo k.

Approach Comparison

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(2^n · n)	O(1)	Enumerates every subset
Optimal Tree DP	O(n · k²)	O(n · k)	Independent-set DP with modulo tracking

Algorithm Walkthrough

DP Definition

For each node u, maintain two arrays of length k.

• dp0[u][r] = number of valid configurations in the subtree of u where:

• node u is not selected

• total sum modulo k equals r

• dp1[u][r] = number of valid configurations in the subtree of u where:

• node u is selected

• total sum modulo k equals r

Initialization

For a node before processing any children:

If u is not selected:

• The subtree contributes an empty set.
• Sum modulo k is 0.

Therefore:

dp0[0] = 1

If u is selected:

• The subset contains only node u.
• Residue is nums[u] % k.

Therefore:

dp1[nums[u] % k] = 1

Processing Children

We process children one at a time and merge their contributions.

Case 1: Current Node Not Selected

If u is not selected, a child may be:

• selected
• not selected

So every child contributes:

childWays = dp0Child + dp1Child

We combine residues using modular convolution.

If the current accumulated residue is a and the child's residue is b, the new residue becomes:

(a + b) % k

Case 2: Current Node Selected

If u is selected, no child may be selected.

Therefore only dp0Child is allowed.

Again we perform residue convolution.

Postorder Traversal

Each node depends on its children, so we must process the tree from leaves upward.

A DFS postorder traversal naturally computes children before parents.

Final Answer

At the root:

answer = dp0Root[0] + dp1Root[0]

This counts all independent subsets whose sum is divisible by k.

However, the empty subset is included in dp0Root[0].

Since the problem asks for non-empty subsets, subtract one:

answer -= 1

Take the result modulo 10^9 + 7.

Why it Works

The DP maintains a complete count of all valid independent subsets inside every subtree.

The state distinguishes whether the current node is selected, which is exactly the information needed to enforce the parent-child restriction.

Every subtree configuration is represented by its sum modulo k, and residues combine correctly through modular addition.

Because every child is merged independently and every valid combination is counted exactly once, the root DP contains all valid independent subsets of the entire tree. Subtracting the unique empty subset leaves precisely the required answer.

Python Solution

from typing import List
import sys

class Solution:
    def countValidSubsets(self, parent: List[int], nums: List[int], k: int) -> int:
        MOD = 10**9 + 7
        n = len(parent)

        children = [[] for _ in range(n)]
        for i in range(1, n):
            children[parent[i]].append(i)

        sys.setrecursionlimit(10000)

        def dfs(u: int):
            dp0 = [0] * k
            dp1 = [0] * k

            dp0[0] = 1
            dp1[nums[u] % k] = 1

            for v in children[u]:
                child0, child1 = dfs(v)

                next0 = [0] * k
                next1 = [0] * k

                # u not selected
                for r1 in range(k):
                    if dp0[r1] == 0:
                        continue
                    for r2 in range(k):
                        ways = (child0[r2] + child1[r2]) % MOD
                        if ways:
                            next0[(r1 + r2) % k] = (
                                next0[(r1 + r2) % k]
                                + dp0[r1] * ways
                            ) % MOD

                # u selected
                for r1 in range(k):
                    if dp1[r1] == 0:
                        continue
                    for r2 in range(k):
                        if child0[r2]:
                            next1[(r1 + r2) % k] = (
                                next1[(r1 + r2) % k]
                                + dp1[r1] * child0[r2]
                            ) % MOD

                dp0 = next0
                dp1 = next1

            return dp0, dp1

        root0, root1 = dfs(0)

        answer = (root0[0] + root1[0] - 1) % MOD
        return answer

Implementation Explanation

The tree is first converted from the parent array into a children adjacency list. This allows efficient traversal from parent to children.

The DFS returns two arrays:

• dp0, representing configurations where the current node is excluded.
• dp1, representing configurations where the current node is included.

Each array stores counts indexed by residue modulo k.

The initialization corresponds to the two possible states of a single node:

• Excluded, producing residue 0.
• Included, producing residue nums[u] % k.

For every child, we create new DP arrays and merge residue distributions through modular convolution.

When the current node is excluded, the child may be either included or excluded. When the current node is included, only the child's excluded state is allowed.

After processing all children, the resulting arrays represent the entire subtree.

At the root, residue 0 corresponds to sums divisible by k. The empty subset contributes exactly one configuration, so we subtract it.

Go Solution

package main

func countValidSubsets(parent []int, nums []int, k int) int {
	const MOD int = 1000000007

	n := len(parent)

	children := make([][]int, n)
	for i := 1; i < n; i++ {
		children[parent[i]] = append(children[parent[i]], i)
	}

	var dfs func(int) ([]int, []int)

	dfs = func(u int) ([]int, []int) {
		dp0 := make([]int, k)
		dp1 := make([]int, k)

		dp0[0] = 1
		dp1[nums[u]%k] = 1

		for _, v := range children[u] {
			child0, child1 := dfs(v)

			next0 := make([]int, k)
			next1 := make([]int, k)

			for r1 := 0; r1 < k; r1++ {
				if dp0[r1] == 0 {
					continue
				}
				for r2 := 0; r2 < k; r2++ {
					ways := (child0[r2] + child1[r2]) % MOD
					if ways == 0 {
						continue
					}
					idx := (r1 + r2) % k
					next0[idx] = (next0[idx] + int((int64(dp0[r1])*int64(ways))%MOD)) % MOD
				}
			}

			for r1 := 0; r1 < k; r1++ {
				if dp1[r1] == 0 {
					continue
				}
				for r2 := 0; r2 < k; r2++ {
					if child0[r2] == 0 {
						continue
					}
					idx := (r1 + r2) % k
					next1[idx] = (next1[idx] + int((int64(dp1[r1])*int64(child0[r2]))%MOD)) % MOD
				}
			}

			dp0 = next0
			dp1 = next1
		}

		return dp0, dp1
	}

	root0, root1 := dfs(0)

	ans := (root0[0] + root1[0] - 1) % MOD
	if ans < 0 {
		ans += MOD
	}

	return ans
}

Go-Specific Notes

Go's int type is large enough to store values modulo 10^9 + 7, but intermediate multiplications can exceed that range. Therefore the implementation performs multiplications using int64 before taking the modulus.

Slices are used for the residue arrays, and each merge allocates fresh arrays of size k.

The recursion depth is at most 1000, which is safe for Go.

Worked Examples

Example 1

parent = [-1,0,1]
nums   = [1,2,3]
k = 3

Tree:

0
|
1
|
2

Node 2:

State	Residue Counts
dp0	[1,0,0]
dp1	[1,0,0]

Since 3 % 3 = 0.

Node 1 initialization:

State	Residue Counts
dp0	[1,0,0]
dp1	[0,0,1]

Merge child 2:

State	Residue Counts
dp0	[2,0,0]
dp1	[0,0,1]

Node 0 initialization:

State	Residue Counts
dp0	[1,0,0]
dp1	[0,1,0]

Merge node 1:

State	Residue Counts
dp0	[2,0,1]
dp1	[0,2,0]

Final:

dp0[0] + dp1[0] = 2

Subtract empty subset:

2 - 1 = 1

Answer:

1

Example 2

parent = [-1,0,0,0]
nums   = [2,1,2,1]
k = 3

Tree:

    0
  / | \
 1  2  3

Leaf DPs:

Node 1:

dp0 = [1,0,0]
dp1 = [0,1,0]

Node 2:

dp0 = [1,0,0]
dp1 = [0,0,1]

Node 3:

dp0 = [1,0,0]
dp1 = [0,1,0]

After merging all three children into the root, the residue count for remainder 0 equals 3.

One of those configurations is the empty subset.

Therefore:

3 - 1 = 2

Answer:

2

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n · k²)	Each tree edge causes one DP merge, each merge performs residue convolution over k × k states
Space	O(n · k)	DP information across recursion and subtree states

The dominant operation is merging residue distributions. Each merge examines all pairs of residues, resulting in O(k²) work. Since the tree has n - 1 edges and k ≤ 100, the overall complexity is efficient for the given constraints.

Test Cases

sol = Solution()

assert sol.countValidSubsets([-1, 0, 1], [1, 2, 3], 3) == 1
# Example 1

assert sol.countValidSubsets([-1, 0, 0, 0], [2, 1, 2, 1], 3) == 2
# Example 2

assert sol.countValidSubsets([-1], [3], 3) == 1
# Single node, divisible

assert sol.countValidSubsets([-1], [2], 3) == 0
# Single node, not divisible

assert sol.countValidSubsets([-1, 0], [1, 2], 3) == 1
# Only child alone forms valid subset

assert sol.countValidSubsets([-1, 0], [3, 3], 3) == 2
# Root alone and child alone

assert sol.countValidSubsets([-1, 0, 0], [1, 1, 1], 2) == 1
# Only {1,2} has even sum

assert sol.countValidSubsets([-1, 0, 1, 2, 3], [1, 1, 1, 1, 1], 2) == 3
# Chain structure

assert sol.countValidSubsets([-1, 0, 0, 0, 0], [5, 5, 5, 5, 5], 5) == 16
# Star tree, every non-empty independent subset of leaves plus root alone

assert sol.countValidSubsets([-1, 0, 0], [10**9, 10**9, 10**9], 1) == 4
# Every non-empty independent subset is valid when k=1

Test Summary

Test	Why
Example 1	Validates sample output
Example 2	Validates sample output
Single divisible node	Smallest valid tree
Single non-divisible node	No valid subset exists
Two-node chain	Parent-child exclusion
Both values divisible	Multiple singleton solutions
Small star tree	Sibling selections allowed
Long chain	Strong adjacency constraints
Large star	Many independent combinations
k = 1	Every sum divisible

Edge Cases

Single Node Tree

When n = 1, the only possible non-empty subset is the root itself. A buggy implementation might accidentally count the empty subset or mishandle the base case. The DP initialization naturally handles this situation because the root starts with both exclusion and inclusion states, and the final subtraction removes only the empty subset.

Long Chain

A chain represents the strongest possible adjacency restriction because every node except the ends has two neighbors. Greedy approaches often fail here because choosing one node immediately affects the next. The DP explicitly tracks whether a node is selected, ensuring that parent-child conflicts are never counted.

Empty Subset Producing Residue Zero

The empty subset always has sum 0, which is divisible by every k. Since the problem requires a non-empty subset, forgetting to remove this configuration would produce an answer that is too large by exactly one whenever residue 0 is considered. The implementation subtracts one from the root's residue-0 count, removing the empty subset and nothing else.

Very Large Node Values

Values may be as large as 10^9. Storing exact sums would be unnecessary and inefficient. The implementation immediately reduces each value using nums[u] % k, because divisibility depends only on the residue modulo k. This keeps the DP state bounded by k ≤ 100.