Rust Interview Puzzles: Partition a vector into 2 slices with equal sum

Rust Interview Puzzles: Partition a vector into 2 slices with equal sum

·

2 min read

Problem

We need to partition a provided vector vec into 2 slices, each having the same sum of the elements.

Example 1:

Input: vec = [7, -5, -4, 2, 4]
Output: ([7, -5], [-4, 2, 4])
Explanation: Each of the 2 partitions has the equal sum of the elements 2.

Example 2:

Input: vec = [7, -6, 3, -5, 1]
Output: ([], [7, -6, 3, -5, 1])
Explanation: Each of the 2 partitions has the equal sum of the elements 0 (the empty partition is assumed to have 0 sum).

Example 3:

Input: vec = [1, 3, 5, 7]
Output: None
Explanation: This vector could not be partitioned (there are no slices with an equal sum).

Solution

fn partition_vector_into_2_slices_with_equal_sum(vec: &Vec<i32>) -> Option<(&[i32], &[i32])> {
    let sum_all = vec.iter().fold(0, |sum, x| sum + x);

    let mut sum_left = 0;
    (0..vec.len()).find_map(|i| {
        let sum_right = sum_all - sum_left;
        let found = if sum_left == sum_right {
            Some((&vec[0..i], &vec[i..vec.len()]))
        } else {
            None
        };
        sum_left += vec[i];
        found
    })
}

Explanation

To solve the puzzle in linear time and without using auxiliary data structures, we could first precalculate the sum of all vector elements sum_all. And then when iterating over the vector vec for the second time, we could maintain sum_left of the elements we have seen so far, and on each iteration check if the current sum_left is equal to sum_right (we could figure out sum_right because we know sum_all and the current sum_left). If that is the case, then we have found 2 partitions (slices) with the equal sum and we wrap them in Some, otherwise we return None for the current iteration. Note that we use find_map function which would return the first Some result (the first pair of found partitions with equal sum) and would ignore None results.

Complexity

The time complexity of this solution is O(n) (n is a size of the input vector), but we had to iterate over the vector vec twice. The auxiliary space complexity is constant (O(1)).


Rust Playground with this code