aboutsummaryrefslogtreecommitdiff
path: root/find_sum_arr_product_seq.py
blob: ebadb9f570873201c13d416eec55aba8caa1f899 (plain)
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
# 3539. Find Sum of Array Product of Magical Sequences

"""
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.len', and the binary representation of '2^seq[0] + 2^seq[1]
+ ...' has 'k' set bits. the array product of this sequence is defined as
'prod(seq) = (nums[seq[0]]  nums[seq[1]]  ...)'. return the sum of the
array products for all valid magical sequences. since the answer may be very
large, return it modulo 10^9+7. a set bit refers to a bit in the binary
represenation of a number that has a value of 1
"""


class Solution(object):
    def magicalSum(self, m, k, nums):
        """
        :type m: int
        :type k: int
        :type nums: List[int]
        :rtype: int
        """
        mod = 10**9 + 7
        n = len(nums)
        powtab = [[1] * (m + 1) for _ in range(n)]
        for i in range(n):
            for c in range(1, m + 1):
                powtab[i][c] = powtab[i][c - 1] * nums[i] % mod
        comb = [[0] * (m + 1) for _ in range(m + 1)]
        for i in range(m + 1):
            comb[i][0] = 1
            for j in range(1, i + 1):
                comb[i][j] = (comb[i - 1][j - 1] + comb[i - 1][j]) % mod
        dp = [[[0] * (k + 1) for _ in range(m + 1)] for __ in range(m + 1)]
        dp[m][0][0] = 1
        for i in range(n):
            next_dp = [[[0] * (k + 1) for _ in range(m + 1)] for __ in range(m + 1)]
            pow_i = powtab[i]
            for rem in range(m + 1):
                for carry in range(m + 1):
                    row = dp[rem][carry]
                    if not any(row_elt for row_elt in row):
                        continue
                    for ones in range(k + 1):
                        base_val = row[ones]
                        if base_val == 0:
                            continue
                        for c in range(rem + 1):
                            t = c + carry
                            bit = t & 1
                            ones2 = ones + bit
                            if ones2 > k:
                                continue
                            carry2 = t >> 1
                            rem2 = rem - c
                            add = base_val
                            add = add * comb[rem][c] % mod
                            add = add * pow_i[c] % mod
                            next_dp[rem2][carry2][ones2] = (
                                next_dp[rem2][carry2][ones2] + add
                            ) % mod
            dp = next_dp
        ans = 0
        for carry in range(m + 1):
            extra = bin(carry).count("1")
            need = k - extra
            if 0 <= need <= k:
                ans = (ans + dp[0][carry][need]) % mod
        return ans % mod


if __name__ == "__main__":
    obj = Solution()
    print(obj.magicalSum(m=5, k=5, nums=[1, 10, 100, 10000, 1000000]))
    print(obj.magicalSum(m=2, k=2, nums=[5, 4, 3, 2, 1]))
    print(obj.magicalSum(m=1, k=1, nums=[28]))