548. Split Array with Equal Sum

Given an integer array nums of length n, return true if there is a triplet (i, j, k) which satisfies the following conditions:

• 0 < i, i + 1 < j, j + 1 < k < n - 1
• The sum of subarrays (0, i - 1), (i + 1, j - 1), (j + 1, k - 1) and (k + 1, n - 1) is equal.

A subarray (l, r) represents a slice of the original array starting from the element indexed l to the element indexed r.

class Solution {
public:
    bool splitArray(vector<int>& A) {
        int n = A.size();
        vector<int> psum(n + 1, 0);
        unordered_map<int, vector<int>> mp;
        for(int i = 0; i < n; i++) {
            mp[psum[i+1] = psum[i] + A[i]].push_back(i);
        }
        
        unordered_set<int> vis[n+1][2];
        
        for(int i = 1; i < n; i++) {
            int sum = psum[i];
            if(i + 1 >= n) continue;
            for(auto& j : mp[psum[i+1] + sum]) {
                if(j + 2 >= n or j <= i) continue;
                if(vis[j][0].count(sum)) continue;
                
                vis[j][0].insert(sum);
                for(auto& k : mp[psum[j+2] + sum]) {
                    if(k + 2 >= n or k <= j + 1) continue;
                    if(vis[k][1].count(sum)) continue;
                    
                    vis[k][1].insert(sum);
                    if(psum.back() == psum[k+2] + sum)
                        return true;
                }
            }
        }
        return false;
    }
};