K-Sum Subarray

  • Find the number of subarrays in an array of i32 that sum to k

The point:

  • prepend the prefix-sum array with 0
  • iterate from 1 because 0 was added

Complexity :

Time Space
O(n²) O(n)
  • O(n²) because we iterate through each value in 2 nested loops. Brute force would be O(n3)
  • O(n) because of space taken by the hash map

About Rust :

fn k_sum_subarrays(nums : &[i32], k:i32) -> usize {
    let n = nums.len();

    // Populate prefix sum array, setting its first element to 0
    let mut prefix_sums = Vec::with_capacity(n+1);
    prefix_sums.push(0);
    let mut current = 0;
    for &num in nums {
        current += num;
        prefix_sums.push(current);
    }
    
    // loop trough valid pairs of prefix_sum values to find all subarray summing to `k`  
    let mut count = 0;
    for j in 1..n+1{
        for i in 1..j+1{
            if prefix_sums[j] - prefix_sums[i-1] == k {
                count += 1 ;
            }
        }
    }
    count
}

fn main(){   // no main() if this code runs in a Jupyter cell 
    println!("{:?}", k_sum_subarrays(&vec![1, 2, -1, 1, 2], 3));    // 3
    println!("{:?}", k_sum_subarrays(&[1, 2, -1, 1, 2], 3));        // 3
} // end of local scope OR end of main()       

Optimization

The point:

  • store encountered prefix sum in a hash map
  • we can then in constant time check if curr_prefix_sum - 1 was encountered

Complexity :

Time Space
O(n) O(n)
  • O(n) because we iterate through each value
  • O(n) because of space taken by the hash map

About Rust :

  • count += *prefix_sum_map.get(&(my_key)).unwrap_or(&0);
  • YES : tested on the Rust Playground
use std::collections::HashMap;

fn k_sum_subarrays_optimized(nums : &[i32], k:i32) -> usize {
    // Initialize the hash map with 0 to handle subarrays
    // that sum to `k` from the start of the array
    let mut prefix_sum_map = HashMap::new();
    prefix_sum_map.insert(0, 1);

    // Populate prefix sum array, setting its first element to 0
    // let mut prefix_sums = Vec::with_capacity(nums.len()+1);
    // prefix_sums.push[0];
    let mut curr_prefix_sum = 0;
    let mut count = 0;
    for &num in nums {
        // update running prefix by adding current number
        curr_prefix_sum += num;

        // If a subarray with sum `k` exists, 
        // increment `count` by the number of times it has been found
        count += *prefix_sum_map.get(&(curr_prefix_sum - k)).unwrap_or(&0);

        // Udate the frequency of `curr_prefix_sum` in the hash map
        let freq = *prefix_sum_map.get(&curr_prefix_sum).unwrap_or(&0);
        prefix_sum_map.insert(curr_prefix_sum, freq+1);
    }
    count
}

fn main(){   // no main() if this code runs in a Jupyter cell 
    println!("{:?}", k_sum_subarrays_optimized(&vec![1, 2, -1, 1, 2], 3));  // 3
    println!("{:?}", k_sum_subarrays_optimized(&[1, 2, -1, 1, 2], 3));      // 3
} // end of local scope OR end of main()       

Back to top

Published on: Jun 22 2025 at 09:00 AM | Last updated: Jun 22 2025 at 09:00 AM

Copyright © 1964-2025 - 40tude