Maximum number of ways to partition an array

Array Hash Table Counting Enumeration Prefix Sum

35.6% acceptance

Feb 25, 2026

520

You are given a 0-indexed integer array nums of length n. You can change at most one element to k. Return the maximum number of pivot indices that satisfy: prefix sum == suffix sum, after at most one change.

Solution

Rust

Time O(n)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn ways_to_partition(nums: Vec<i32>, k: i32) -> i32 {
    use std::collections::HashMap;
    let n = nums.len();
    let nums: Vec<i64> = nums.iter().map(|&x| x as i64).collect();
    let k = k as i64;
    
    // prefix[i] = sum of nums[0..=i]
    let mut prefix = vec![0i64; n];
    prefix[0] = nums[0];
    for i in 1..n { prefix[i] = prefix[i-1] + nums[i]; }
    let total = prefix[n-1];
    
    // Case 1: no change. Count i in [1,n-1] where prefix[i-1] == total - prefix[i-1]
    // i.e., 2*prefix[i-1] == total
    let no_change = (1..n).filter(|&i| 2 * prefix[i-1] == total).count() as i32;
    
    // Case 2: change nums[j] to k, diff = k - nums[j]
    // New total = total + diff
    // Need 2*prefix[i-1] + (if i-1 >= j then diff else 0) == total + diff
    // If j < i: 2*(prefix[i-1]+diff) == total+diff => 2*prefix[i-1] == total - diff
    // If j >= i: 2*prefix[i-1] == total + diff
    
    // Build frequency maps: prefix_count_left[v] = count of prefix[0..j-1]
    // As j increases from 0 to n-1, we want:
    // count of prefix[i-1] (for i from 1 to j) where 2*prefix[i-1] == total+diff (j >= i means j can be i..n-1, so these are i <= j)
    // count of prefix[i-1] (for i from j+1 to n-1) where 2*prefix[i-1] == total-diff (j < i)
    
    // Precompute suffix count map
    let mut suffix_map: HashMap<i64, i64> = HashMap::new();
    for i in 0..n-1 { *suffix_map.entry(prefix[i]).or_insert(0) += 1; }
    
    let mut prefix_map: HashMap<i64, i64> = HashMap::new();
    let mut best = no_change;
    
    for j in 0..n {
      let diff = k - nums[j];
      let new_total = total + diff;
      
      // i <= j: need 2*prefix[i-1] == new_total (prefix[i-1] counted in prefix_map, keys = prefix[0..j-1])
      // i > j: need 2*prefix[i-1] == total - diff + 0... wait:
      // prefix becomes prefix[i-1] + diff for i-1 >= j, i.e. i >= j+1
      // new_prefix[i-1] = prefix[i-1] + diff if i-1 >= j else prefix[i-1]
      // For pivot at i: new_prefix[i-1] == new_total - new_prefix[i-1]
      // If i-1 < j (i.e. i <= j): new_prefix[i-1] = prefix[i-1], need 2*prefix[i-1] == new_total
      // If i-1 >= j (i.e. i >= j+1): new_prefix[i-1] = prefix[i-1]+diff, need 2*(prefix[i-1]+diff) == new_total
      //   => 2*prefix[i-1] == new_total - 2*diff = total - diff
      
      let cnt = *prefix_map.get(&(new_total - new_total%2*0)).unwrap_or(&0); // dummy, fix below
      
      // Count from prefix_map: i <= j means keys prefix[0..j-1], want 2*v == new_total
      let cnt1 = if new_total % 2 == 0 { *prefix_map.get(&(new_total/2)).unwrap_or(&0) } else { 0 };
      // Count from suffix_map: keys prefix[j..n-2], want 2*v == total - diff
      let want2 = total - diff;
      let cnt2 = if want2 % 2 == 0 { *suffix_map.get(&(want2/2)).unwrap_or(&0) } else { 0 };
      
      let _ = cnt; // suppress warning
      best = best.max((cnt1 + cnt2) as i32);
      
      // Move prefix[j] from suffix to prefix
      if j < n-1 {
        let v = prefix[j];
        *suffix_map.entry(v).or_insert(0) -= 1;
        *prefix_map.entry(v).or_insert(0) += 1;
      }
    }
    
    best
  }
}