548. Split Array with Equal Sum

Given an array with n integers, you need to find if there are triplets (i, j, k) which satisfies following conditions:

  1. 0 < i, i + 1 < j, j + 1 < k < n - 1
  2. Sum of subarrays (0, i - 1), (i + 1, j - 1), (j + 1, k - 1) and (k + 1, n - 1) should be equal.
where we define that subarray (L, R) represents a slice of the original array starting from the element indexed L to the element indexed R.

Example:

Input: [1,2,1,2,1,2,1]
Output: True
Explanation:
i = 1, j = 3, k = 5. 
sum(0, i - 1) = sum(0, 0) = 1
sum(i + 1, j - 1) = sum(2, 2) = 1
sum(j + 1, k - 1) = sum(4, 4) = 1
sum(k + 1, n - 1) = sum(6, 6) = 1

Note:
  1. 1 <= n <= 2000.
  2. Elements in the given array will be in range [-1,000,000, 1,000,000].

Rust Solution

struct Solution;
use std::collections::HashSet;

impl Solution {
    fn split_array(nums: Vec<i32>) -> bool {
        let n = nums.len();
        let mut prefix = vec![0; n];
        let mut prev = 0;
        if n < 7 {
            return false;
        }
        for i in 0..n {
            prev += nums[i];
            prefix[i] = prev;
        }
        for j in 3..n - 3 {
            let mut hs: HashSet<i32> = HashSet::new();
            for i in 1..j - 1 {
                let a = prefix[i - 1];
                let b = prefix[j - 1] - prefix[i];
                if a == b {
                    hs.insert(a);
                }
            }
            for k in j + 1..n - 1 {
                let c = prefix[k - 1] - prefix[j];
                let d = prefix[n - 1] - prefix[k];
                if c == d {
                    if hs.contains(&c) {
                        return true;
                    }
                }
            }
        }
        false
    }
}

#[test]
fn test() {
    let nums = vec![1, 2, 1, 2, 1, 2, 1];
    let res = true;
    assert_eq!(Solution::split_array(nums), res);
}

Having problems with this solution? Click here to submit an issue on github.