Count Stable Segments in an Array with Prefix Sums

 avatar
unknown
c_cpp
a year ago
521 B
15
Indexable
using ll = long long;
int countStableSegments(vector<int> arr) {
    int n = arr.size();
    if(n <= 2)
        return 0;
    unordered_map<int, unordered_map<ll, ll>>dp;//[cur][psum]
    dp[arr[0]][arr[0]] = 1;
    ll cur_sum = arr[0], ans = 0;
    for(int i = 1; i < n; i++){
        int x = arr[i];
        if(dp.count(arr[i])){
            if(dp[x].count(cur_sum-x))
                ans += dp[x][cur_sum-x];
        }
        cur_sum += arr[i];
        dp[arr[i]][cur_sum]++;
    }
    return ans;
}
Editor is loading...
Leave a Comment