Farmer John Solves 3SUM

Explicación

Aunque el problema puede parecer intimidante al principio, ayuda pensar en cómo un rango más pequeño está conectado a otro.

Si $\texttt{ways}[i][j]$ es igual al número de tripletas no ordenadas que suman cero entre $[i, j]$ , ¿de qué estados depende esto?

Podemos elegir excluir un extremo o el otro, lo que nos da una imagen como esta (recordatorio de no contar dos veces la superposición).

Aunque esto contará todas las tripletas entre ellos, no estamos teniendo en cuenta ninguna tripleta que use ambos elementos en $i$ y $j$ , ya que esas no están incluidas en nuestra transición “tipo PIE”.^{Nuestra transición toma todas las tripletas de $[i + 1, j]$ y $[i, j - 1]$}

Para hacer esto, podemos precalcular el número de tripletas que empiezan en $i$ y terminan en $j$ en tiempo $\mathcal{O}(N^2)$ almacenando las ocurrencias del complemento de $i$ y $j$ en un arreglo.

Más formalmente, nuestra transición final es:

\texttt{ways}[i][j] = \texttt{ways}[i + 1][j] + \texttt{ways}[i][j - 1] - \texttt{ways}[i + 1][j - 1] + \texttt{trp}[i][j]

donde $\texttt{trp}[i][k]$ es igual al número de tripletas que suman cero y empiezan en $i$ y terminan en $k$ .

Implementación

Complejidad temporal: $\mathcal{O}(N^2)$

Nótese que no podemos almacenar $\texttt{trp}$ en su propio arreglo debido a las restricciones de memoria.

C++Java


#include <bits/stdc++.h>
using namespace std;
 
using ll = long long;
const int MAX_VAL = 1e6;
 
 
int main() {
	freopen("threesum.in", "r", stdin);
	freopen("threesum.out", "w", stdout);
	cin.tie(0)->sync_with_stdio(0);
 
	int n, q;
	cin >> n >> q;
 
	vector<int> val(n);
	for (int i = 0; i < n; i++) {
		cin >> val[i];
		val[i] += MAX_VAL;
	}
 
	// number of triplets such that i and k are fixed
	vector<vector<ll>> ways(n, vector<ll>(n, 0));
	vector<int> frq(2 * MAX_VAL);
 
	for (int i = n - 2; i >= 0; i--) {
		// be careful to only consider values in between i and j
		frq[val[i + 1]]++;
 
		for (int j = i + 2; j < n; j++) {
			// separate val from shifted amt
			int compliment = (MAX_VAL * 3) - (val[i] + val[j]);
 
			if (compliment >= 0 && compliment < (2 * MAX_VAL)) {
				ways[i][j] = frq[compliment];
			}
			frq[val[j]]++;
		}
 
		for (int j = i + 1; j < n; j++) { frq[val[j]]--; }
	}
 
	// essentially 2D prefix sum on ways
	for (int i = n - 1; i >= 0; i--) {
		for (int j = i + 1; j < n; j++) {
			ways[i][j] += (ways[i + 1][j] + ways[i][j - 1] - ways[i + 1][j - 1]);
		}
	}
 
	for (int _ = 0; _ < q; _++) {
		int l, r;
		cin >> l >> r;
		cout << ways[l - 1][r - 1] << '\n';
	}
}


import java.io.*;
import java.util.*;
public class Threesum {
	static final int MAX_VAL = 1000000;
	public static void main(String[] args) throws IOException {
		Kattio io = new Kattio("threesum");
		int N = io.nextInt();
		int Q = io.nextInt();
		int[] a = new int[N];
		for (int i = 0; i < N; i++) { a[i] = io.nextInt(); }
		long[][] dp = new long[N + 1][N + 1];
		int[] cnt = new int[2 * MAX_VAL + 1];
		// Calculate the number of k such that a[i] + a[j] + a[k] = 0 given i
		// and j
		for (int i = 0; i < N - 1; i++) {
			for (int j = 0; j < N; j++) { cnt[a[j] + MAX_VAL] = 0; }
			for (int j = i + 1; j < N; j++) {
				int k = -a[i] - a[j];
				if (k >= -MAX_VALUE && k <= MAX_VALUE) {
					dp[i + 1][j + 1] += cnt[k + MAX_VAL];
				}
				cnt[a[j] + MAX_VAL]++;
			}
		}
 
		for (int i = N; i >= 1; i--) {
			for (int j = i + 1; j <= N; j++) {
				dp[i][j] += dp[i + 1][j] + dp[i][j - 1] - dp[i + 1][j - 1];
			}
		}
		for (int i = 0; i < Q; i++) {
			int a1 = io.nextInt();
			int b = io.nextInt();
			io.println(dp[a1][b]);
		}
		io.close();
	}
 
	// CodeSnip{Kattio}
}