1398C - Good Subarrays - CodeForces Solution


data structures dp math *1600

Please click on ads to support us..

Python Code:

from math import ceil
import os
import sys
from io import BytesIO, IOBase
 
 
BUFSIZE: int = 8192
 
 
class FastIO(IOBase):
    newlines = 0
 
    def __init__(self, file):
        self._fd = file.fileno()
        self.buffer = BytesIO()
        self.writable = "x" in file.mode or "r" not in file.mode
        self.write = self.buffer.write if self.writable else None
 
    def read(self):
        while True:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            if not b:
                break
            ptr = self.buffer.tell()
            self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr)
        self.newlines = 0
        return self.buffer.read()
 
    def readline(self):
        while self.newlines == 0:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            self.newlines = b.count(b"\n") + (not b)
            ptr = self.buffer.tell()
            self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr)
        self.newlines -= 1
        return self.buffer.readline()
 
    def flush(self):
        if self.writable:
            os.write(self._fd, self.buffer.getvalue())
            self.buffer.truncate(0), self.buffer.seek(0)
 
 
class IOWrapper(IOBase):
    def __init__(self, file):
        self.buffer = FastIO(file)
        self.flush = self.buffer.flush
        self.writable = self.buffer.writable
        self.write = lambda s: self.buffer.write(s.encode("ascii"))
        self.read = lambda: self.buffer.read().decode("ascii")
        self.readline = lambda: self.buffer.readline().decode("ascii")
 
 
def print(*args, **kwargs):
    
    sep, file = kwargs.pop("sep", " "), kwargs.pop("file", sys.stdout)
    at_start = True
    for x in args:
        if not at_start:
            file.write(sep)
        file.write(str(x))
        at_start = False
    file.write(kwargs.pop("end", "\n"))
    if kwargs.pop("flush", False):
        file.flush()
 
 
sys.stdin, sys.stdout = IOWrapper(sys.stdin), IOWrapper(sys.stdout)
input = lambda: sys.stdin.readline().rstrip("\r\n")
 
 
if __name__ == "__main__":
    pass
def inpi():
    return(int(input()))
def inpl():
    return(list(map(int,input().split())))
def inps():
    s = input()
    return(list(s[:len(s)]))
def inpm():
    return(map(int,input().split()))
def inpst():
    return(set(map(int,input().split())))


    
def solve():
    n = inpi()
    a = [int(i)for i in input()]
    pre = [0]*n
    d = {0:1}
    res = 0
    for i in range(n):
        pre[i] = pre[i-1] + a[i]
                if pre[i] - i - 1 in d:
            res += (d[pre[i] - i - 1])
            d[pre[i] - i - 1] += 1
        else:
            d[pre[i] - i - 1] = 1
        print(res)
    
        
'''

'''

    
t = 1
t = inpi()
for _ in range(t):
    solve()


				 		 	 	 					 	 		 		   		

C++ Code:

#include <bits/stdc++.h>
//#include <ext/pb_ds/assoc_container.hpp>
//#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
//using namespace __gnu_pbds;
#define U unsigned
#define I long long int
#define S string
#define C char
#define D long double
#define A auto
#define B bool
#define pb push_back
#define mp make_pair
#define ff first 
#define ss second
#define all(x) (x).begin(),(x).end()
#define debug(x) cout<<#x<<" is "<<x<<"\n"
#define printp(x); for(A n:x){cout<<n.ff<<" "<<n.ss<<"\n";}
#define printa(x); for(A n:x){cout<<n<<" ";}cout<<"\n";
#define V(x) vector<x>
#define W(t) int t;cin>>t;while(t--)
#define asc(i,x,n) for(I i=x;i<n;i++)
#define dsc(i,x,n) for(I i=x;i>=n;i--)
#define mod 1000000007
#define endl "\n"
#define yes cout << "YES\n";
#define no cout << "NO\n";
#define google cout << "Case #" << i+1 << ": ";
#define Yellow_Flash ios_base::sync_with_stdio(false);cin.tie(NULL);cout.tie(NULL); init_code();
void init_code(){
  #ifndef ONLINE_JUDGE
  freopen("input.txt", "r", stdin);
  freopen("output.txt", "w", stdout);
  #endif
}
//remember modinv
void solve(){
  I n;
  cin>>n;
  S s;
  cin>>s;
  V(I) v;
  v.pb(0);
  I t=0;
  for (I i = 0; i < n; ++i){
    v.pb(t+s[i]-'0'-i-1);
    t+=s[i]-'0';
  }
  map<I,I> m;
  for(A x:v){
    // cout<<x<<" ";
    m[x]++;
  }
  // cout<<endl;
  t=0;
  for(A x:m){
    t+=x.ss*(x.ss-1);
  }
  t/=2;
  cout<<t<<endl;
}
int main(){
  Yellow_Flash
  I t=1;
  cin>>t; 
  for (I i = 0; i < t; ++i){
    // google
    solve();   
   }
  return 0;
}
//Number of digits in N = floor(log10(N)) + 1;
//Number is power of 2 : return x && (!(x&(x-1)));


Comments

Submit
0 Comments
More Questions

1588. Sum of All Odd Length Subarrays
1662. Check If Two String Arrays are Equivalent
1832. Check if the Sentence Is Pangram
1678. Goal Parser Interpretation
1389. Create Target Array in the Given Order
1313. Decompress Run-Length Encoded List
1281. Subtract the Product and Sum of Digits of an Integer
1342. Number of Steps to Reduce a Number to Zero
1528. Shuffle String
1365. How Many Numbers Are Smaller Than the Current Number
771. Jewels and Stones
1512. Number of Good Pairs
672. Richest Customer Wealth
1470. Shuffle the Array
1431. Kids With the Greatest Number of Candies
1480. Running Sum of 1d Array
682. Baseball Game
496. Next Greater Element I
232. Implement Queue using Stacks
844. Backspace String Compare
20. Valid Parentheses
746. Min Cost Climbing Stairs
392. Is Subsequence
70. Climbing Stairs
53. Maximum Subarray
1527A. And Then There Were K
1689. Partitioning Into Minimum Number Of Deci-Binary Numbers
318. Maximum Product of Word Lengths
448. Find All Numbers Disappeared in an Array
1155. Number of Dice Rolls With Target Sum