#3539
Hard Algorithms Find sum of array product of magical sequences
Array Math Dynamic Programming Bit Manipulation Combinatorics Bitmask
62.0% acceptance
Feb 25, 2026
217
156
You are given two integers m and k, and an integer array nums.
A sequence seq of size m is magical if 0 <= seq[i] < nums.length and the binary
representation of 2^seq[0] + 2^seq[1] + ... + 2^seq[m-1] has exactly k set bits.
The array product is nums[seq[0]] * nums[seq[1]] * ... * nums[seq[m-1]].
Return the sum of array products for all valid magical sequences, modulo 10^9+7.
Solution
Rust
Time O(n * m)
Space O(n * m)
impl Solution {
pub fn magical_sum(m: i32, k: i32, nums: Vec<i32>) -> i32 {
// A magical sequence of length m chooses m indices (with repetition allowed).
// Their sum 2^a0 + 2^a1 + ... + 2^a(m-1) has exactly k set bits.
//
// Key insight: think of how many times each index i appears in the sequence.
// Let cnt[i] = # of times index i is chosen. Then sum(cnt) = m.
// The sum 2^a0+...+2^a(m-1) = sum_i cnt[i] * 2^i (in binary, carry-propagated).
// We need this to have exactly k set bits.
//
// The contribution to array product is multinomial:
// C(m; cnt[0], cnt[1], ...) * prod_i nums[i]^cnt[i]
//
// But with nums.len() up to 50 and m up to 30, we use carry-aware DP.
//
// Process indices 0..N one by one. For each index i, choose how many times
// to pick it (cnt_i). Track the carry when adding cnt_i * 2^i to the running sum.
// The carry represents the bits above the current position.
//
// dp[carry][bits_set][items_placed] = sum of (multinomial * product contributions)
//
// At each index position i, we process carry + cnt_i:
// new_bit = (carry + cnt_i) % 2 -> contributes to bit i
// new_carry = (carry + cnt_i) / 2
// if new_bit == 1, bits_set increases by 1
//
// We track remaining items to place (m - placed so far).
// At the end (after all N indices), carry must be 0 and bits_set == k.
//
// State: dp[carry][bits][placed] but carry can grow...
// max carry: at each step carry <= m/2... at most m = 30, so carry <= 30.
// After processing all n indices, remaining carry must propagate into higher bits.
const MOD: i64 = 1_000_000_007;
let n = nums.len();
let m = m as usize;
let k = k as usize;
// Precompute multinomial factorials
let mut fact = vec![1i64; m + 1];
for i in 1..=m { fact[i] = fact[i - 1] * i as i64 % MOD; }
let mut inv_fact = vec![1i64; m + 1];
inv_fact[m] = mod_pow(fact[m], (MOD - 2) as u64, MOD);
for i in (0..m).rev() { inv_fact[i] = inv_fact[i + 1] * (i + 1) as i64 % MOD; }
// dp[carry][bits_set][placed] = weighted count
// carry up to m = 30, bits_set up to k <= 30, placed up to m = 30
let max_carry = m + 1;
let mut dp = vec![vec![vec![0i64; m + 1]; k + 1]; max_carry];
dp[0][0][0] = 1;
for _i in 0..n {
let nv = nums[_i] as i64 % MOD;
let mut ndp = vec![vec![vec![0i64; m + 1]; k + 1]; max_carry];
for carry in 0..max_carry {
for bits in 0..=k {
for placed in 0..=m {
let val = dp[carry][bits][placed];
if val == 0 { continue; }
let remaining = m - placed;
// choose cnt in 0..=remaining for this index
for cnt in 0..=remaining {
let total = carry + cnt;
let new_bit = total % 2;
let new_carry = total / 2;
let new_bits = bits + new_bit;
if new_bits > k { continue; }
if new_carry >= max_carry { continue; }
let new_placed = placed + cnt;
// contribution: C(remaining, cnt) * nv^cnt
let coeff = fact[remaining] % MOD * inv_fact[cnt] % MOD
* inv_fact[remaining - cnt] % MOD
* mod_pow(nv, cnt as u64, MOD) % MOD;
ndp[new_carry][new_bits][new_placed] =
(ndp[new_carry][new_bits][new_placed] + val * coeff) % MOD;
}
}
}
}
dp = ndp;
}
// After all indices, propagate remaining carry
// carry bits propagate: bit at position n, n+1, ...
// We need total bits_set == k for the final number
// Carry itself contributes additional set bits (popcount of carry)
let mut ans = 0i64;
for carry in 0..max_carry {
for bits in 0..=k {
let val = dp[carry][bits][m];
if val == 0 { continue; }
let extra_bits = carry.count_ones() as usize;
if bits + extra_bits == k {
ans = (ans + val) % MOD;
}
}
}
ans as i32
}
}
fn mod_pow(mut base: i64, mut exp: u64, modulus: i64) -> i64 {
let mut result = 1i64;
base %= modulus;
while exp > 0 {
if exp & 1 == 1 { result = result * base % modulus; }
exp >>= 1;
base = base * base % modulus;
}
result
}