LeetCode 3539 - Find Sum of Array Product of Magical Sequences

LeetCode Problem 3539

Difficulty: 🔴 Hard
Topics: Array, Math, Dynamic Programming, Bit Manipulation, Combinatorics, Bitmask

Solution

Problem Understanding

The expression in the statement should be interpreted as:

$$2^{seq[0]} + 2^{seq[1]} + \cdots + 2^{seq[m-1]}$$

A sequence seq of length m is magical when the binary representation of this sum contains exactly k set bits.

Each element of the sequence is an index into nums. The contribution of an index i is the power of two $2^i$. If an index appears multiple times, the corresponding power of two is added multiple times, which may create binary carries.

For every magical sequence, we compute

$$prod(seq)=\prod nums[seq[i]]$$

and we must return the sum of these products over all magical sequences.

The constraints are the key observation:

• m ≤ 30
• k ≤ 30
• nums.length ≤ 50

The number of possible sequences is

$$(\text{nums.length})^m$$

which is astronomically large. Therefore, we must count them indirectly using combinatorics and dynamic programming.

An important detail is that the condition depends only on how many times each index appears, not on the order of appearances. However, the product contribution of a count distribution must still include all ordered sequences that realize that distribution.

Several edge cases are important:

• Multiple occurrences of the same index create binary carries.
• The final carry after processing the largest index contributes additional set bits.
• nums.length can be much larger than m, so many indices may never be chosen.
• k may be as small as 1 or as large as m.

Approaches

Brute Force

A brute force solution would generate every sequence of length m.

For each sequence:

1. Compute the sum of powers of two.
2. Count its set bits.
3. If the count equals k, multiply the corresponding values from nums and add the result.

This is correct because it directly follows the definition of a magical sequence.

Unfortunately, the number of sequences is

$$n^m$$

where $n = nums.length$.

With $n=50$ and $m=30$, this is completely infeasible.

Key Insight

Instead of thinking about ordered sequences, think about how many times each index is used.

Let:

$$c_i$$

be the number of times index i appears.

Then:

$$\sum_i c_i = m$$

and the power-of-two sum becomes

$$\sum_i c_i 2^i.$$

The number of ordered sequences corresponding to a count vector is the multinomial coefficient

$$\frac{m!}{\prod_i c_i!}.$$

The product contribution is

$$\prod_i nums[i]^{c_i}.$$

Therefore each count vector contributes

$$m!\prod_i \frac{nums[i]^{c_i}}{c_i!}.$$

Now observe that the binary popcount of

$$\sum_i c_i 2^i$$

can be computed by processing bits from low to high while carrying overflow exactly as binary addition does.

Since m ≤ 30, every count and carry is at most 30, making a digit-DP over positions, carries, used elements, and current popcount feasible.

Approach Comparison

Approach	Time Complexity	Space Complexity	Notes
Brute Force	$O(n^m \cdot m)$	$O(m)$	Enumerates every sequence
Optimal	$O(n \cdot m^2 \cdot k \cdot m)$	$O(n \cdot m \cdot k \cdot m)$ via memoization	DP on counts and binary carries

Algorithm Walkthrough

1. Let n = len(nums).
2. Precompute factorials and inverse factorials up to m.

The contribution of choosing index i exactly c times is

$$\frac{nums[i]^c}{c!}.$$

We precompute these values for every index and every possible count. 3. Define a memoized DP state:

$$dfs(pos,\ remaining,\ carry,\ ones)$$

where:

• pos is the current index in nums
• remaining is how many sequence positions still need to be assigned
• carry is the incoming binary carry
• ones is the number of set bits already produced

4. At position pos, choose how many times this index appears:

$$c \in [0, remaining]$$ 5. The total value accumulated at this bit position is

$$carry + c.$$

Therefore:

• current bit

$$bit=(carry+c)\bmod 2$$

• next carry

$$nextCarry=\lfloor (carry+c)/2 \rfloor$$ 6. Add the weight

$$\frac{nums[pos]^c}{c!}$$

and recurse. 7. After all indices have been processed:

• all selections must be used (remaining == 0)
• the remaining carry still contains higher binary digits

The final popcount becomes

$$ones + popcount(carry).$$ 8. Accept the state only if this equals k. 9. The DP computes

$$\prod_i \frac{nums[i]^{c_i}}{c_i!}.$$

Multiply the final result by m! to restore the multinomial coefficient.

Why it works

The DP enumerates every possible count vector $(c_0,c_1,\ldots,c_{n-1})$ satisfying $\sum c_i=m$. For each count vector, binary carries are propagated exactly as they would be when evaluating $\sum c_i2^i$, so the computed popcount is correct. The weight accumulated by the DP is $\prod nums[i]^{c_i}/c_i!$, and multiplying by $m!$ converts it into the multinomial count of all ordered sequences realizing that count vector. Thus every magical sequence contributes exactly once and every non-magical sequence contributes zero.

Python Solution

from typing import List
from functools import lru_cache

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

        fact = [1] * (m + 1)
        for i in range(1, m + 1):
            fact[i] = fact[i - 1] * i % MOD

        inv_fact = [1] * (m + 1)
        inv_fact[m] = pow(fact[m], MOD - 2, MOD)
        for i in range(m, 0, -1):
            inv_fact[i - 1] = inv_fact[i] * i % MOD

        weight = [[0] * (m + 1) for _ in range(n)]
        for i, val in enumerate(nums):
            power = 1
            for c in range(m + 1):
                weight[i][c] = power * inv_fact[c] % MOD
                power = power * val % MOD

        @lru_cache(None)
        def dfs(pos: int, remaining: int, carry: int, ones: int) -> int:
            if ones > k:
                return 0

            if pos == n:
                if remaining != 0:
                    return 0
                return 1 if ones + carry.bit_count() == k else 0

            res = 0

            for cnt in range(remaining + 1):
                total = carry + cnt
                bit = total & 1
                next_carry = total >> 1

                res += (
                    weight[pos][cnt]
                    * dfs(
                        pos + 1,
                        remaining - cnt,
                        next_carry,
                        ones + bit,
                    )
                )
                res %= MOD

            return res

        return fact[m] * dfs(0, m, 0, 0) % MOD

The implementation follows the mathematical derivation directly.

The factorial and inverse factorial arrays are used to encode the multinomial coefficient implicitly. The weight table stores

$$nums[i]^c / c!$$

for every index and count.

The memoized DFS processes indices one at a time. For each possible count assigned to the current index, it updates the binary carry and the number of set bits generated at that position.

Once all indices have been processed, the remaining carry contains all higher-order bits. Their popcount is added to the accumulated count, and the state is accepted only when it equals k.

Finally, the answer is multiplied by m!, converting the normalized weights back into the correct multinomial contributions.

Go Solution

package main

func magicalSum(m int, k int, nums []int) int {
	const MOD int = 1000000007

	n := len(nums)

	fact := make([]int, m+1)
	invFact := make([]int, m+1)

	fact[0] = 1
	for i := 1; i <= m; i++ {
		fact[i] = int(int64(fact[i-1]) * int64(i) % MOD)
	}

	powMod := func(a, e int) int {
		res := 1
		x := a % MOD
		for e > 0 {
			if e&1 == 1 {
				res = int(int64(res) * int64(x) % MOD)
			}
			x = int(int64(x) * int64(x) % MOD)
			e >>= 1
		}
		return res
	}

	invFact[m] = powMod(fact[m], MOD-2)
	for i := m; i >= 1; i-- {
		invFact[i-1] = int(int64(invFact[i]) * int64(i) % MOD)
	}

	weight := make([][]int, n)
	for i := 0; i < n; i++ {
		weight[i] = make([]int, m+1)

		p := 1
		for c := 0; c <= m; c++ {
			weight[i][c] = int(int64(p) * int64(invFact[c]) % MOD)
			p = int(int64(p) * int64(nums[i]) % MOD)
		}
	}

	type State struct {
		pos       int
		remaining int
		carry     int
		ones      int
	}

	memo := make(map[State]int)

	var dfs func(int, int, int, int) int

	dfs = func(pos, remaining, carry, ones int) int {
		if ones > k {
			return 0
		}

		if pos == n {
			if remaining != 0 {
				return 0
			}

			if ones+popcount(carry) == k {
				return 1
			}
			return 0
		}

		st := State{pos, remaining, carry, ones}
		if v, ok := memo[st]; ok {
			return v
		}

		res := 0

		for cnt := 0; cnt <= remaining; cnt++ {
			total := carry + cnt
			bit := total & 1
			nextCarry := total >> 1

			cur := int(
				int64(weight[pos][cnt]) *
					int64(dfs(
						pos+1,
						remaining-cnt,
						nextCarry,
						ones+bit,
					)) % MOD,
			)

			res += cur
			if res >= MOD {
				res -= MOD
			}
		}

		memo[st] = res
		return res
	}

	ans := int(int64(fact[m]) * int64(dfs(0, m, 0, 0)) % MOD)
	return ans
}

func popcount(x int) int {
	cnt := 0
	for x > 0 {
		cnt++
		x &= x - 1
	}
	return cnt
}

The Go implementation mirrors the Python solution closely. The main differences are the explicit modular exponentiation function, a struct-based memoization key, and careful use of int64 during multiplication to avoid overflow before taking the modulus.

Worked Examples

Example 1

Input:

m = 5
k = 5
nums = [1,10,100,10000,1000000]

The only way to obtain exactly five set bits using five selections is to choose each exponent once:

Index	Count
0	1
1	1
2	1
3	1
4	1

Then

$$2^0+2^1+2^2+2^3+2^4=31$$

and

$$popcount(31)=5.$$

The product is

$$1\cdot10\cdot100\cdot10000\cdot1000000=10^{13}.$$

There are

$$5!=120$$

permutations of these indices.

Therefore:

$$120 \cdot 10^{13}$$

modulo $10^9+7$ equals:

991600007

Example 2

Input:

m = 2
k = 2
nums = [5,4,3,2,1]

A pair contributes two set bits exactly when the indices are different.

Examples:

Sequence	Binary Sum	Popcount
[0,0]	2	1
[0,1]	3	2
[1,3]	10	2
[4,4]	32	1

Thus all ordered pairs with distinct indices are magical.

The total is

$$\sum_{i\neq j} nums[i]nums[j].$$

Using

$$(\sum nums)^2-\sum nums_i^2$$

gives:

$$15^2-55=170.$$

Answer:

170

Example 3

Input:

m = 1
k = 1
nums = [28]

The only sequence is:

[0]

Its power sum is:

$$2^0=1$$

which has one set bit.

Its product is:

$$28.$$

Answer:

28

Complexity Analysis

Measure	Complexity	Explanation
Time	$O(n \cdot m^3 \cdot k)$	DP states are bounded by position, remaining count, carry, and popcount
Space	$O(n \cdot m^2 \cdot k)$	Memoization stores all reachable states

The crucial observation is that both the carry and the total number of selected elements are at most 30. This keeps the state space small enough despite nums.length reaching 50.

Test Cases

s = Solution()

assert s.magicalSum(5, 5, [1, 10, 100, 10000, 1000000]) == 991600007  # example 1
assert s.magicalSum(2, 2, [5, 4, 3, 2, 1]) == 170  # example 2
assert s.magicalSum(1, 1, [28]) == 28  # example 3

assert s.magicalSum(1, 1, [7, 9, 11]) == 27  # every single choice valid

assert s.magicalSum(2, 1, [1]) == 1  # repeated index creates carry

assert s.magicalSum(2, 2, [1, 1]) == 2  # two ordered distinct pairs

assert s.magicalSum(3, 1, [2]) == 8  # only one index, heavy carry chain

assert s.magicalSum(3, 3, [1, 1, 1]) > 0  # maximal popcount case

assert s.magicalSum(30, 1, [1]) == 1  # largest m with single index

assert s.magicalSum(4, 2, [3, 5]) >= 0  # mixed carry behavior

Test	Why
Example 1	Official sample
Example 2	Official sample
Example 3	Official sample
m=1 with multiple indices	Every index independently valid
Single index, m=2	Tests carry generation
Two indices, equal values	Tests ordered counting
Single index, m=3	Tests repeated carry propagation
k=m style scenario	Tests maximum popcount target
m=30, one index	Largest allowed m
Two-index mixed case	General carry handling

Edge Cases

Repeated Use of the Same Index

If an index is chosen multiple times, the resulting power of two is added repeatedly. For example, choosing index 0 twice gives:

$$2^0+2^0=2.$$

The result has only one set bit because a carry occurs. A naive solution that simply counts distinct indices would fail. The DP explicitly propagates carries at every bit position, so this case is handled correctly.

Remaining Carry After the Last Index

After processing all indices, there may still be a nonzero carry. For example:

$$30 \cdot 2^0.$$

The binary representation extends beyond the highest index present in nums. Ignoring this carry would produce an incorrect popcount. The implementation adds carry.bit_count() when the recursion reaches the end.

Large nums.length with Small m

The array can contain up to 50 indices while only 30 selections are made. Most indices will often receive count zero. The DP naturally handles this by allowing cnt = 0 at every position, ensuring all valid count distributions are considered without wasting work on impossible selections.