Skip to content

3539. Find Sum of Array Product of Magical Sequences

Description

You are given two integers, m and k, and an integer array nums.

A sequence of integers seq is called magical if:

  • seq has a size of m.
  • 0 <= seq[i] < nums.length
  • The binary representation of 2seq[0] + 2seq[1] + ... + 2seq[m - 1] has k set bits.

The array product of this sequence is defined as prod(seq) = (nums[seq[0]] * nums[seq[1]] * ... * nums[seq[m - 1]]).

Return the sum of the array products for all valid magical sequences.

Since the answer may be large, return it modulo 109 + 7.

A set bit refers to a bit in the binary representation of a number that has a value of 1.

 

Example 1:

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

Output: 991600007

Explanation:

All permutations of [0, 1, 2, 3, 4] are magical sequences, each with an array product of 1013.

Example 2:

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

Output: 170

Explanation:

The magical sequences are [0, 1], [0, 2], [0, 3], [0, 4], [1, 0], [1, 2], [1, 3], [1, 4], [2, 0], [2, 1], [2, 3], [2, 4], [3, 0], [3, 1], [3, 2], [3, 4], [4, 0], [4, 1], [4, 2], and [4, 3].

Example 3:

Input: m = 1, k = 1, nums = [28]

Output: 28

Explanation:

The only magical sequence is [0].

 

Constraints:

  • 1 <= k <= m <= 30
  • 1 <= nums.length <= 50
  • 1 <= nums[i] <= 108

Solutions

We design a function \(\text{dfs}(i, j, k, st)\), which represents the number of ways when we are currently processing the \(i\)-th element of array \(\textit{nums}\), still need to select numbers from the remaining \(j\) positions to fill into the magical sequence, still need to satisfy having \(k\) set bits in binary form, and the current carry from the previous bit is \(st\). Then the answer is \(\text{dfs}(0, m, k, 0)\).

The execution process of function \(\text{dfs}(i, j, k, st)\) is as follows:

If \(k < 0\) or \(i = n\) and \(j > 0\), it means the current solution is not feasible, return \(0\).

If \(i = n\), it means we have finished processing array \(\textit{nums}\). We need to check if there are still set bits in the current carry \(st\), and if so, we need to decrease \(k\). If \(k = 0\) at this point, it means the current solution is feasible, return \(1\), otherwise return \(0\).

Otherwise, we enumerate selecting \(t\) numbers at position \(i\) to fill into the magical sequence (\(0 \leq t \leq j\)). The number of ways to fill \(t\) numbers into the magical sequence is \(\binom{j}{t}\), the array product is \(\textit{nums}[i]^t\), the updated carry is \((t + st) >> 1\), the updated number of required set bits is \(k - ((t + st) \& 1)\), and we recursively call \(\text{dfs}(i + 1, j - t, k - ((t + st) \& 1), (t + st) >> 1)\). The sum of all \(t\) solutions is \(\text{dfs}(i, j, k, st)\).

To efficiently compute the binomial coefficient \(\binom{m}{n}\), we preprocess the factorial array \(f\) and the inverse factorial array \(g\), where \(f[i] = i! \mod (10^9 + 7)\) and \(g[i] = (i!)^{-1} \mod (10^9 + 7)\). Then \(\binom{m}{n} = f[m] \cdot g[n] \cdot g[m - n] \mod (10^9 + 7)\).

The time complexity is \(O(n \cdot m^3 \cdot k)\) and the space complexity is \(O(n \cdot m^2 \cdot k)\), where \(n\) is the length of array \(\textit{nums}\), and \(m\) and \(k\) are the parameters in the problem.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
mx = 30
mod = 10**9 + 7
f = [1] + [0] * mx
g = [1] + [0] * mx

for i in range(1, mx + 1):
    f[i] = f[i - 1] * i % mod
    g[i] = pow(f[i], mod - 2, mod)


def comb(m: int, n: int) -> int:
    return f[m] * g[n] * g[m - n] % mod


class Solution:
    def magicalSum(self, m: int, k: int, nums: List[int]) -> int:
        @cache
        def dfs(i: int, j: int, k: int, st: int) -> int:
            if k < 0 or (i == len(nums) and j > 0):
                return 0
            if i == len(nums):
                while st:
                    k -= st & 1
                    st >>= 1
                return int(k == 0)
            res = 0
            for t in range(j + 1):
                nt = t + st
                p = pow(nums[i], t, mod)
                nk = k - (nt & 1)
                res += comb(j, t) * p * dfs(i + 1, j - t, nk, nt >> 1)
                res %= mod
            return res

        ans = dfs(0, m, k, 0)
        dfs.cache_clear()
        return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
class Solution {
    static final int N = 31;
    static final long MOD = 1_000_000_007L;
    private static final long[] f = new long[N];
    private static final long[] g = new long[N];
    private Long[][][][] dp;

    static {
        f[0] = 1;
        g[0] = 1;
        for (int i = 1; i < N; ++i) {
            f[i] = f[i - 1] * i % MOD;
            g[i] = qpow(f[i], MOD - 2);
        }
    }

    public static long qpow(long a, long k) {
        long res = 1;
        while (k != 0) {
            if ((k & 1) == 1) {
                res = res * a % MOD;
            }
            a = a * a % MOD;
            k >>= 1;
        }
        return res;
    }

    public static long comb(int m, int n) {
        return f[m] * g[n] % MOD * g[m - n] % MOD;
    }

    public int magicalSum(int m, int k, int[] nums) {
        int n = nums.length;
        dp = new Long[n + 1][m + 1][k + 1][N];
        long ans = dfs(0, m, k, 0, nums);
        return (int) ans;
    }

    private long dfs(int i, int j, int k, int st, int[] nums) {
        if (k < 0 || (i == nums.length && j > 0)) {
            return 0;
        }
        if (i == nums.length) {
            while (st > 0) {
                k -= (st & 1);
                st >>= 1;
            }
            return k == 0 ? 1 : 0;
        }

        if (dp[i][j][k][st] != null) {
            return dp[i][j][k][st];
        }

        long res = 0;
        for (int t = 0; t <= j; t++) {
            int nt = t + st;
            int nk = k - (nt & 1);
            long p = qpow(nums[i], t);
            long tmp = comb(j, t) * p % MOD * dfs(i + 1, j - t, nk, nt >> 1, nums) % MOD;
            res = (res + tmp) % MOD;
        }

        return dp[i][j][k][st] = res;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
const int N = 31;
const long long MOD = 1'000'000'007;

long long f[N], g[N];

long long qpow(long long a, long long k) {
    long long res = 1;
    while (k) {
        if (k & 1) res = res * a % MOD;
        a = a * a % MOD;
        k >>= 1;
    }
    return res;
}

int init = []() {
    f[0] = g[0] = 1;
    for (int i = 1; i < N; ++i) {
        f[i] = f[i - 1] * i % MOD;
        g[i] = qpow(f[i], MOD - 2);
    }
    return 0;
}();

long long comb(int m, int n) {
    return f[m] * g[n] % MOD * g[m - n] % MOD;
}

class Solution {
    vector<vector<vector<vector<long long>>>> dp;

    long long dfs(int i, int j, int k, int st) {
        if (k < 0 || (i == nums.size() && j > 0)) {
            return 0;
        }
        if (i == nums.size()) {
            while (st > 0) {
                k -= (st & 1);
                st >>= 1;
            }
            return k == 0 ? 1 : 0;
        }

        long long& res = dp[i][j][k][st];
        if (res != -1) {
            return res;
        }

        res = 0;
        for (int t = 0; t <= j; ++t) {
            int nt = t + st;
            int nk = k - (nt & 1);
            long long p = qpow(nums[i], t);
            long long tmp = comb(j, t) * p % MOD * dfs(i + 1, j - t, nk, nt >> 1) % MOD;
            res = (res + tmp) % MOD;
        }
        return res;
    }

public:
    int magicalSum(int m, int k, vector<int>& nums) {
        int n = nums.size();
        this->nums = nums;
        dp.assign(n + 1, vector<vector<vector<long long>>>(m + 1, vector<vector<long long>>(k + 1, vector<long long>(N, -1))));
        return dfs(0, m, k, 0);
    }

private:
    vector<int> nums;
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
const N = 31
const MOD = 1_000_000_007

var f [N]int64
var g [N]int64

func init() {
    f[0], g[0] = 1, 1
    for i := 1; i < N; i++ {
        f[i] = f[i-1] * int64(i) % MOD
        g[i] = qpow(f[i], MOD-2)
    }
}

func qpow(a, k int64) int64 {
    res := int64(1)
    for k > 0 {
        if k&1 == 1 {
            res = res * a % MOD
        }
        a = a * a % MOD
        k >>= 1
    }
    return res
}

func comb(m, n int) int64 {
    if n < 0 || n > m {
        return 0
    }
    return f[m] * g[n] % MOD * g[m-n] % MOD
}

func magicalSum(m int, k int, nums []int) int {
    n := len(nums)
    dp := make([][][][]int64, n+1)
    for i := 0; i <= n; i++ {
        dp[i] = make([][][]int64, m+1)
        for j := 0; j <= m; j++ {
            dp[i][j] = make([][]int64, k+1)
            for l := 0; l <= k; l++ {
                dp[i][j][l] = make([]int64, N)
                for s := 0; s < N; s++ {
                    dp[i][j][l][s] = -1
                }
            }
        }
    }

    var dfs func(i, j, k, st int) int64
    dfs = func(i, j, k, st int) int64 {
        if k < 0 || (i == n && j > 0) {
            return 0
        }
        if i == n {
            for st > 0 {
                k -= st & 1
                st >>= 1
            }
            if k == 0 {
                return 1
            }
            return 0
        }
        if dp[i][j][k][st] != -1 {
            return dp[i][j][k][st]
        }
        res := int64(0)
        for t := 0; t <= j; t++ {
            nt := t + st
            nk := k - (nt & 1)
            p := qpow(int64(nums[i]), int64(t))
            tmp := comb(j, t) * p % MOD * dfs(i+1, j-t, nk, nt>>1) % MOD
            res = (res + tmp) % MOD
        }
        dp[i][j][k][st] = res
        return res
    }

    return int(dfs(0, m, k, 0))
}

Comments