Hard
题目描述
给定两个整数 m 和 k,以及一个整数数组 nums。
如果一个整数序列 seq 满足以下条件,则称其为神奇的:
seq的大小为m0 <= seq[i] < nums.length2^seq[0] + 2^seq[1] + ... + 2^seq[m-1]的二进制表示中有k个设置位
该序列的数组乘积定义为 prod(seq) = (nums[seq[0]] * nums[seq[1]] * ... * nums[seq[m-1]])。
返回所有有效神奇序列的数组乘积之和。
由于答案可能很大,请返回结果对 10^9 + 7 取模的值。
设置位是指数字二进制表示中值为 1 的位。
示例 1:
输入:m = 5, k = 5, nums = [1,10,100,10000,1000000]
输出:991600007
解释:[0, 1, 2, 3, 4] 的所有排列都是神奇序列,每个的数组乘积都是 10^13。
示例 2:
输入:m = 2, k = 2, nums = [5,4,3,2,1]
输出:170
解释:神奇序列有 [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], [4, 3]。
示例 3:
输入:m = 1, k = 1, nums = [28]
输出:28
解释:唯一的神奇序列是 [0]。
约束:
1 <= k <= m <= 301 <= nums.length <= 501 <= nums[i] <= 10^8
解题思路
这是一道结合了动态规划、位运算和组合数学的复杂题目。
核心思路:
题目要求找出所有长度为 m 的序列,使得 2^seq[0] + 2^seq[1] + ... + 2^seq[m-1] 的二进制表示恰好有 k 个1。
关键观察:当我们计算多个2的幂次的和时,会产生进位现象。例如 2^1 + 2^1 = 2^2,这样原本应该有2个1位的和变成了只有1个1位。
动态规划状态设计:
使用 dp[i][j][mask] 表示:
i:已选择的数字个数j:当前和的二进制表示中的1位个数mask:低位的进位状态(表示哪些位上有未处理的进位)
对于每个状态,我们需要跟踪两个值:
- 达到该状态的方案数
- 达到该状态的所有乘积之和
状态转移:
对于每个可选的索引,我们需要:
- 计算新的进位状态
- 更新1位的个数
- 累加方案数和乘积和
组合数学优化:
当数组中有重复元素时,我们可以用组合数学来优化计算,避免重复枚举相同的排列。
最终答案是所有满足条件状态(选择了 m 个数且恰好有 k 个1位)的乘积和。
代码实现
class Solution {
public:
int magicalSum(int m, int k, vector<int>& nums) {
const int MOD = 1e9 + 7;
int n = nums.size();
// dp[i][j][mask] = {count, sum}
map<tuple<int, int, int>, pair<long long, long long>> dp;
dp[{0, 0, 0}] = {1, 1};
for (int i = 0; i < m; i++) {
map<tuple<int, int, int>, pair<long long, long long>> next_dp;
for (auto& [state, val] : dp) {
auto [pos, bits, mask] = state;
auto [count, sum] = val;
for (int j = 0; j < n; j++) {
int new_mask = mask | (1 << j);
int new_bits = bits;
// Calculate carries
int carry_mask = mask & (1 << j) ? mask : 0;
if (mask & (1 << j)) {
new_mask &= ~(1 << j);
carry_mask = mask;
} else {
new_bits++;
}
// Process carries
while (carry_mask) {
int next_carry = 0;
for (int bit = 0; bit < 30; bit++) {
if (carry_mask & (1 << bit)) {
if (new_mask & (1 << (bit + 1))) {
next_carry |= (1 << (bit + 1));
} else {
new_mask |= (1 << (bit + 1));
new_bits++;
}
new_mask &= ~(1 << bit);
new_bits--;
}
}
carry_mask = next_carry;
}
auto new_state = make_tuple(i + 1, new_bits, new_mask);
long long new_count = (count * 1) % MOD;
long long new_sum = (sum * nums[j]) % MOD;
if (next_dp.find(new_state) != next_dp.end()) {
next_dp[new_state].first = (next_dp[new_state].first + new_count) % MOD;
next_dp[new_state].second = (next_dp[new_state].second + new_sum) % MOD;
} else {
next_dp[new_state] = {new_count, new_sum};
}
}
}
dp = next_dp;
}
long long result = 0;
for (auto& [state, val] : dp) {
auto [pos, bits, mask] = state;
if (pos == m && bits == k && mask == 0) {
result = (result + val.second) % MOD;
}
}
return result;
}
};
class Solution:
def magicalSum(self, m: int, k: int, nums: List[int]) -> int:
MOD = 10**9 + 7
n = len(nums)
# dp[i][j][mask] = (count, sum)
dp = {(0, 0, 0): (1, 1)}
for i in range(m):
next_dp = {}
for (pos, bits, mask), (count, total_sum) in dp.items():
for j in range(n):
new_mask = mask
new_bits = bits
if mask & (1 << j):
# There's already a carry at position j
new_mask &= ~(1 << j) # Remove the bit
# Process carry
carry_pos = j
while carry_pos < 30:
if new_mask & (1 << (carry_pos + 1)):
new_mask &= ~(1 << (carry_pos + 1))
carry_pos += 1
else:
new_mask |= (1 << (carry_pos + 1))
break
else:
# No carry, just set the bit
new_mask |= (1 << j)
new_bits += 1
# Count set bits in final mask
final_bits = new_bits
temp_mask = new_mask
while temp_mask:
if temp_mask & 1:
final_bits += 1
temp_mask >>= 1
# Only consider states with reasonable bit counts
if final_bits <= k:
new_state = (i + 1, new_bits, new_mask)
new_count = count % MOD
new_sum = (total_sum * nums[j]) % MOD
if new_state in next_dp:
next_dp[new_state] = (
(next_dp[new_state][0] + new_count) % MOD,
(next_dp[new_state][1] + new_sum) % MOD
)
else:
next_dp[new_state] = (new_count, new_sum)
dp = next_dp
result = 0
for (pos, bits, mask), (count, total_sum) in dp.items():
if pos == m and bits == k and mask == 0:
result = (result + total_sum) % MOD
return result
public class Solution {
public int MagicalSum(int m, int k, int[] nums) {
const int MOD = 1000000007;
int n = nums.Length;
var dp = new Dictionary<(int, int, int), (long, long)>();
dp[(0, 0, 0)] = (1, 1);
for (int i = 0; i < m; i++) {
var nextDp = new Dictionary<(int, int, int), (long, long)>();
foreach (var kvp in dp) {
var (pos, bits, mask) = kvp.Key;
var (count, sum) = kvp.Value;
for (int j = 0; j < n; j++) {
int newMask = mask;
int newBits = bits;
if ((mask & (1 << j)) != 0) {
newMask &= ~(1 << j);
int carryPos = j;
while (carryPos < 30) {
if ((newMask & (1 << (carryPos + 1))) != 0) {
newMask &= ~(1 << (carryPos + 1));
carryPos++;
} else {
newMask |= (1 << (carryPos + 1));
break;
}
}
} else {
newMask |= (1 << j);
newBits++;
}
int finalBits = newBits;
int tempMask = newMask;
while (tempMask > 0) {
if ((tempMask & 1) == 1) {
finalBits++;
}
tempMask >>= 1;
}
if (finalBits <= k) {
var newState = (i + 1, newBits, newMask);
long newCount = count % MOD;
long newSum = (sum * nums[j]) % MOD;
if (nextDp.ContainsKey(newState)) {
var existing = nextDp[newState];
nextDp[newState] = (
(existing.Item1 + newCount) % MOD,
(existing.Item2 + newSum) % MOD
);
} else {
nextDp[newState] = (newCount, newSum);
}
}
}
}
dp = nextDp;
}
long result = 0;
foreach (var kvp in dp) {
var (pos, bits, mask) = kvp.Key;
if (pos == m && bits == k && mask == 0) {
result = (result + kvp.Value.Item2) % MOD;
}
}
return (int)result;
}
}
var magicalSum = function(m, k, nums) {
const MOD = 1000000007n;
const n = nums.length;
const combinations = Array.from(
{length: m + 1},
() => new Array(m + 1).fill(0n)
);
combinations[0][0] = 1n;
for (let total = 1; total <= m; total++) {
combinations[total][0] = 1n;
combinations[total][total] = 1n;
for (let chosen = 1; chosen < total; chosen++) {
combinations[total][chosen] =
(combinations[total - 1][chosen - 1] +
combinations[total - 1][chosen]) % MOD;
}
}
const weights = nums.map(value => {
const powers = new Array(m + 1).fill(1n);
const base = BigInt(value) % MOD;
for (let count = 1; count <= m; count++) {
powers[count] = powers[count - 1] * base % MOD;
}
return Array.from({length: m + 1}, (_, used) => {
const result = new Array(m - used + 1);
for (let count = 0; count + used <= m; count++) {
result[count] = combinations[used + count][count] * powers[count] % MOD;
}
return result;
});
});
const encode = (used, ones, carry) =>
(used * (k + 1) + ones) * (m + 1) + carry;
const stateCount = (m + 1) * (k + 1) * (m + 1);
let dp = new BigUint64Array(stateCount);
dp[encode(0, 0, 0)] = 1n;
for (let index = 0; index < n; index++) {
const next = new BigUint64Array(stateCount);
for (let used = 0; used <= m; used++) {
for (let ones = 0; ones <= k; ones++) {
for (let carry = 0; carry <= m; carry++) {
const value = dp[encode(used, ones, carry)];
if (value === 0n) continue;
for (let count = 0; count + used <= m; count++) {
const sum = carry + count;
const nextOnes = ones + (sum & 1);
if (nextOnes > k) continue;
const nextState = encode(used + count, nextOnes, sum >> 1);
const contribution = value * weights[index][used][count] % MOD;
next[nextState] = (next[nextState] + contribution) % MOD;
}
}
}
}
dp = next;
}
const popcount = value => {
let count = 0;
while (value > 0) {
count += value & 1;
value >>= 1;
}
return count;
};
let result = 0n;
for (let ones = 0; ones <= k; ones++) {
for (let carry = 0; carry <= m; carry++) {
if (ones + popcount(carry) === k) {
result = (result + dp[encode(m, ones, carry)]) % MOD;
}
}
}
return Number(result);
};
复杂度分析
| 指标 | 复杂度 |
|---|---|
| 时间 | - |
| 空间 | - |