CF - Good Subarrays

Authors: Jesse Choe, Brad Ma

Edit This Page

Official Editorial (Python)

Solution - Prefix Sums + Math

Following the editorial, we build a prefix sum array pp on the existing array.

We know that subarray formed by [l,r)[l, r) is a good subarray iff r−l=pr−plr-l=p_r-p_l. Rearranging this equation leads to pr−r=pl−lp_r-r=p_l-l, so we build a map (sum_dist) on values of pi−ip_i-i for all valid ii.

Then, we iterate over all values of the map and check how many unordered pairs we can build with the number of values of ii that have the same value of pi−ip_i-i.

Implementation

Time Complexity: O(Nlog⁡N)\mathcal{O}(N\log N)

C++

#include <bits/stdc++.h>
using namespace std;


using ll = long long;


void solve() {
	int n;
	cin >> n;
	vector<int> pref_arr(n + 1);
	for (int i = 1; i <= n; i++) {

Java

import java.io.*;
import java.util.*;


public class GoodSubarrays {
	static long solve(int arrLen, String strArr) {
		int[] prefArr = new int[arrLen + 1];
		for (int i = 1; i <= arrLen; i++) { prefArr[i] = strArr.charAt(i - 1) - '0'; }
		for (int i = 1; i <= arrLen; i++) { prefArr[i] += prefArr[i - 1]; }


		Map<Integer, Long> sumDist = new HashMap<>();

Python

from collections import defaultdict


for _ in range(int(input())):
	arr_len = int(input())
	arr = [ord(i) - ord("0") for i in input()]


	pref_arr = [0] + arr
	for i in range(1, len(pref_arr)):
		pref_arr[i] += pref_arr[i - 1]

Join the USACO Forum!

Stuck on a problem, or don't understand a module? Join the USACO Forum and get help from other competitive programmers!

Join Forum