Add solution and test-cases for problem 3539

Author: 0xff-dev

leetcode/3501-3600/3539.Find-Sum-of-Array-Product-of-Magical-Sequences/README.md

@@ -1,28 +1,57 @@
 # [3539.Find Sum of Array Product of Magical Sequences][title]
 
-> [!WARNING|style:flat]
-> This question is temporarily unanswered if you have good ideas. Welcome to [Create Pull Request PR](https://github.com/kylesliu/awesome-golang-algorithm)
-
 ## 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 `2^seq[0] + 2^seq[1] + ... + 2^seq[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** `10^9 + 7`.
+
+A **set bit** refers to a bit in the binary representation of a number that has a value of 1.
 
 **Example 1:**
 
 ```
-Input: a = "11", b = "1"
-Output: "100"
+nput: 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:
 
-### 思路1
-> ...
-Find Sum of Array Product of Magical Sequences
-```go
+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].
+```
 
 ## 结语
 

leetcode/3501-3600/3539.Find-Sum-of-Array-Product-of-Magical-Sequences/Solution.go

@@ -1,5 +1,84 @@
 package Solution
 
-func Solution(x bool) bool {
-	return x
+import "math/bits"
+
+func quickmul(x, y, mod int64) int64 {
+	res, cur := int64(1), x%mod
+	for y > 0 {
+		if y&1 == 1 {
+			res = res * cur % mod
+		}
+		y >>= 1
+		cur = cur * cur % mod
+	}
+	return res
+}
+
+func Solution(m int, k int, nums []int) int {
+	mod := int64(1000000007)
+	fac := make([]int64, m+1)
+	fac[0] = 1
+	for i := 1; i <= m; i++ {
+		fac[i] = fac[i-1] * int64(i) % mod
+	}
+	ifac := make([]int64, m+1)
+	ifac[0] = 1
+	ifac[1] = 1
+	for i := 2; i <= m; i++ {
+		ifac[i] = quickmul(int64(i), mod-2, mod)
+	}
+	for i := 2; i <= m; i++ {
+		ifac[i] = ifac[i-1] * ifac[i] % mod
+	}
+
+	numsPower := make([][]int64, len(nums))
+	for i := range nums {
+		numsPower[i] = make([]int64, m+1)
+		numsPower[i][0] = 1
+		for j := 1; j <= m; j++ {
+			numsPower[i][j] = numsPower[i][j-1] * int64(nums[i]) % mod
+		}
+	}
+
+	f := make([][][][]int64, len(nums))
+	for i := range nums {
+		f[i] = make([][][]int64, m+1)
+		for j := 0; j <= m; j++ {
+			f[i][j] = make([][]int64, m*2+1)
+			for p := 0; p <= m*2; p++ {
+				f[i][j][p] = make([]int64, k+1)
+			}
+		}
+	}
+
+	for j := 0; j <= m; j++ {
+		f[0][j][j][0] = numsPower[0][j] * ifac[j] % mod
+	}
+	for i := 0; i+1 < len(nums); i++ {
+		for j := 0; j <= m; j++ {
+			for p := 0; p <= m*2; p++ {
+				for q := 0; q <= k; q++ {
+					q2 := p%2 + q
+					if q2 > k {
+						break
+					}
+					for r := 0; r+j <= m; r++ {
+						p2 := p/2 + r
+						f[i+1][j+r][p2][q2] += f[i][j][p][q] * numsPower[i+1][r] % mod * ifac[r] % mod
+						f[i+1][j+r][p2][q2] %= mod
+					}
+				}
+			}
+		}
+	}
+
+	res := int64(0)
+	for p := 0; p <= m*2; p++ {
+		for q := 0; q <= k; q++ {
+			if bits.OnesCount(uint(p))+q == k {
+				res = (res + f[len(nums)-1][m][p][q]*fac[m]%mod) % mod
+			}
+		}
+	}
+	return int(res)
 }

leetcode/3501-3600/3539.Find-Sum-of-Array-Product-of-Magical-Sequences/Solution_test.go

@@ -10,30 +10,31 @@ func TestSolution(t *testing.T) {
 	//	测试用例
 	cases := []struct {
 		name   string
-		inputs bool
-		expect bool
+		m, k   int
+		nums   []int
+		expect int
 	}{
-		{"TestCase", true, true},
-		{"TestCase", true, true},
-		{"TestCase", false, false},
+		{"TestCase1", 5, 5, []int{1, 10, 100, 10000, 1000000}, 991600007},
+		{"TestCase2", 2, 2, []int{5, 4, 3, 2, 1}, 170},
+		{"TestCase3", 1, 1, []int{28}, 28},
 	}
 
 	//	开始测试
 	for i, c := range cases {
 		t.Run(c.name+" "+strconv.Itoa(i), func(t *testing.T) {
-			got := Solution(c.inputs)
+			got := Solution(c.m, c.k, c.nums)
 			if !reflect.DeepEqual(got, c.expect) {
-				t.Fatalf("expected: %v, but got: %v, with inputs: %v",
-					c.expect, got, c.inputs)
+				t.Fatalf("expected: %v, but got: %v, with inputs: %v %v %v",
+					c.expect, got, c.m, c.k, c.nums)
 			}
 		})
 	}
 }
 
-//	压力测试
+// 压力测试
 func BenchmarkSolution(b *testing.B) {
 }
 
-//	使用案列
+// 使用案列
 func ExampleSolution() {
 }