E. Count the Subarrays. | Binary Search & Algorithms Practice Problems

Programming tutorials, Coding problems, and Practice questions how to solve this ?

2 Likes

The sum of subarray is just the difference between 2 prefix sums. Since we only care about the absolute value, the actual order of the prefix sums doesn’t matter, and we can sort them and find how many pairs have a difference larger than k by using 2 pointers.

Implementation
#include <iostream>
#include <bits/stdc++.h>
#define ll long long int
using namespace std;
void solve(){
    int n;
    ll k;
    cin>>n>>k;
    vector<ll> prefix(n+1);
    for(int i=1;i<=n;i++){
        int x;
        cin>>x;
        prefix[i]=prefix[i-1]+x;
    }
    sort(prefix.begin(), prefix.end());
    ll ans=0;
    for(int i=0,j=-1;i<=n;i++){
        while(j<n && prefix[i]-prefix[j+1]>k){
            ++j;
        }
        ans+=j+1;
    }
    cout<<ans<<"\n";
}
int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    int t;
    cin>>t;
    while(t--){
        solve();
    }
}
Alternative solution(Not recommended)

At each step of the prefix sum, you can maintain an indexed set that contains all the prefix sums that are before the current index. Now we just need to find the number of values lesser than prefixsum[i]-k and greater than prefixsum[i]+k.

implementation
#include <iostream>
#include <bits/stdc++.h>
#include<ext/pb_ds/assoc_container.hpp>
#include<ext/pb_ds/tree_policy.hpp>
#define ll long long int
#define mp make_pair 
using namespace std;
using namespace __gnu_pbds;
using indexed_set= tree<pair<ll,int>, null_type, less<pair<ll,int>>, rb_tree_tag, tree_order_statistics_node_update>;
void solve(){
    int n;
    ll k;
    cin>>n>>k;
    vector<ll> prefix(n+1);
    for(int i=1;i<=n;i++){
        int x;
        cin>>x;
        prefix[i]=prefix[i-1]+x;
    }
    indexed_set curr;
    ll ans=0;
    for(int i=1;i<=n;i++){
        curr.insert(mp(prefix[i-1], i-1));
        ans+=curr.order_of_key(mp(prefix[i]-k, -1));
        ans+=curr.size() - curr.order_of_key(mp(prefix[i]+k, n+1));
    }
    cout<<ans<<'\n';
}
int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    int t;
    cin>>t;
    while(t--){
        solve();
    }
}
1 Like

thanx

@everule1 will it be valid for (sum is equal to k case) ?
if not how to think for that ?