Soma

Solution by TioPatinhas — 100 pts

Code / Notes

#include <bits/stdc++.h>
using namespace std;
int main() {
    int n;
    long long k;
    cin >> n >> k;
    long long res = 0;
    vector<int>A(n);
    for (int i = 0; i < n; i++) {
        cin >> A[i];
    }
    if (k == 0) {
        long long zeros = 0;
        for (int i = 0; i < n; i++) {
            if(A[i] == 0) {
                zeros++;
                res += zeros;
            } else {
                zeros = 0;
            }
        }
        cout << res;
        return 0;
    }
    long long soma = 0;
    int esq = 0;
    int dir = 0;
    while (dir < n) {
        soma += A[dir];
        while(soma > k && esq <= dir) {
            soma -= A[esq];
            esq++;
        }
        if(soma == k) {
            long long zerosprox = 0;
            int dir2 = dir;
            while (dir2 + 1 < n && A[dir2+1] == 0) {
                zerosprox++;
                dir2++;
            }
            int esq2 = esq;
            long long zerosesq = 0;
            while (esq2 <= dir && A[esq2] == 0) {
                zerosesq++;
                esq2++;
            }
            res += (1 + zerosesq) * (1 + zerosprox);
            esq = esq2;
            if (esq <= dir) { 
            soma -= A[esq];
            esq++;
            }
        }
        dir++;
    }
    cout << res << endl;
}

Last updated 2 weeks, 3 days ago

« Back to problem