Split array with same average

Array Hash Table Math Dynamic Programming Bit Manipulation Bitmask

26.8% acceptance

Feb 22, 2026

1342

142

You are given an integer array nums. You should move each element of nums into one of the two arrays A and B such that A and B are non-empty, and average(A) == average(B). Return true if it is possible to achieve that and false otherwise. Note that for an array arr, average(arr) is the sum of all the elements of arr over the length of arr.

Solution

Rust

Time O(n²)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn split_array_same_average(nums: Vec<i32>) -> bool {
    let n = nums.len();
    let total: i32 = nums.iter().sum();
    // We need k elements summing to total*k/n for some 1<=k<n
    // Use DP: dp[k][s] = can we pick k elements summing to s
    // Optimize: only check k up to n/2
    let half = n / 2;
    // dp[s] = set of possible k values for sum s
    // Better: dp as bitmask over k in 0..=half
    let max_sum = (total as usize).min(nums.iter().map(|&x| x as usize).sum());
    // dp[s] = bitmask of achievable k's
    let mut dp = vec![0u32; max_sum + 1];
    dp[0] = 1u32; // k=0, sum=0
    for &x in &nums {
      let x = x as usize;
      for s in (x..=max_sum).rev() {
        dp[s] |= dp[s - x] << 1;
      }
    }
    for k in 1..=half {
      let s = total as usize * k;
      if s % n == 0 && (dp[s / n] >> k) & 1 == 1 {
        return true;
      }
    }
    false

  }
}