Skip to Content
CTRL K

Good Subarrays

Editorial oficial (Python) 

Solución - Sumas de prefijos + matemáticas

Siguiendo la editorial, construimos un arreglo de sumas de prefijos pp sobre el arreglo existente.

Sabemos que el subarreglo formado por [l,r)[l, r) es un buen subarreglo sii r−l=pr−plr-l=p_r-p_l. Reordenar esta ecuación lleva a pr−r=pl−lp_r-r=p_l-l, así que construimos un mapa (sum_dist) sobre los valores de pi−ip_i-i para todo ii válido.

Luego iteramos sobre todos los valores del mapa y comprobamos cuántos pares no ordenados podemos formar con el número de valores de ii que tienen el mismo valor de pi−ip_i-i.

Implementación

Complejidad temporal: O(Nlog⁡N)\mathcal{O}(N\log N)

#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++) {
		char c;
		cin >> c;
		pref_arr[i] = c - '0';
	}
 
	for (int i = 1; i <= n; i++) { pref_arr[i] += pref_arr[i - 1]; }
 
	map<int, ll> sum_dist;
	for (int i = 0; i <= n; i++) { sum_dist[pref_arr[i] - i]++; }
 
	ll good_arrays = 0;
	for (const auto &[_, f] : sum_dist) {
		// calcular el # de pares no ordenados posibles con f valores de i
		good_arrays += f * (f - 1) / 2;
	}
 
	cout << good_arrays << endl;
}
 
int main() {
	int t;
	cin >> t;
	for (int i = 0; i < t; i++) { solve(); }
}
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<>();
		for (int i = 0; i <= arrLen; i++) {
			int val = prefArr[i] - i;
			sumDist.put(val, sumDist.getOrDefault(val, 0L) + 1);
		}
 
		long goodArrays = 0;
		for (long f : sumDist.values()) {
			// calcular el # de pares no ordenados posibles con f valores de i
			goodArrays += f * (f - 1) / 2;
		}
		return goodArrays;
	}
 
	public static void main(String[] args) {
		Kattio io = new Kattio();
		int t = io.nextInt();
		for (int i = 0; i < t; i++) {
			int arrLen = io.nextInt();
			String array = io.next();
			io.println(solve(arrLen, array));
		}
		io.close();
	}
 
	// CodeSnip{Kattio}
}
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]
 
	sum_dist = defaultdict(int)
	for i in range(len(pref_arr)):
		sum_dist[pref_arr[i] - i] += 1
 
	good_arrays = 0
	for f in sum_dist.values():
		# calcular el # de pares no ordenados posibles con f valores de i
		good_arrays += f * (f - 1) // 2
	print(good_arrays)
Fundamentos