Skip to Content
CTRL K

K-Inversions

Explicación

Queremos el número de índices ii que cumplen:

si=Bandsi+k=A s_i = B \quad \text{and} \quad s_{i+k} = A

Definimos dos arreglos:

  • Bi=1B_i = 1 si si=Bs_i = B, y 00 en caso contrario
  • Ai=1A_i = 1 si si=As_i = A, y 00 en caso contrario

Entonces el número de kk-inversiones es igual a:

∑i=0n−1Bi⋅Ai+k \sum_{i=0}^{n-1} B_i \cdot A_{i+k}

Convertir a multiplicación de polinomios

Para computar esto con una convolución, primero revertimos un arreglo (revertamos B).

Ahora construimos dos polinomios:

P(x)=∑i=0n−1Aixi P(x) = \sum_{i=0}^{n-1} A_i x^i Q(x)=∑i=0n−1Bixi Q(x) = \sum_{i=0}^{n-1} B_i x^i

Sea C(x)=P(x)⋅Q(x)C(x) = P(x) \cdot Q(x). Entonces el coeficiente C[n - 1 + k] es exactamente el número de k-inversiones.

Así, el problema entero se reduce a calcular una sola multiplicación de polinomios y extraer coeficientes específicos.

Optimizar con FFT

Hacerlo por fuerza bruta corre en tiempo O(n2)\mathcal{O}(n^2).

Con la Transformada Rápida de Fourier (FFT), la convolución corre en tiempo O(nlog⁡n)\mathcal{O}(n \log n).

Implementación

Complejidad temporal: O(nlog⁡n)\mathcal{O}(n \log n)

 
#include <bits/stdc++.h>
using namespace std;
 
using cd = complex<double>;
const double PI = acos(-1);
 
// BeginCodeSnip{FFT Template}
void fft(vector<cd> &a, bool inv) {
	int n = a.size();
	if (n == 1) return;
	vector<cd> a0(n / 2), a1(n / 2);
	for (int i = 0; 2 * i < n; i++) a0[i] = a[2 * i], a1[i] = a[2 * i + 1];
	fft(a0, inv);
	fft(a1, inv);
	double ang = 2 * PI / n * (inv ? -1 : 1);
	cd w(1), wn(cos(ang), sin(ang));
	for (int i = 0; 2 * i < n; i++) {
		a[i] = a0[i] + w * a1[i];
		a[i + n / 2] = a0[i] - w * a1[i];
		if (inv) a[i] /= 2, a[i + n / 2] /= 2;
		w *= wn;
	}
}
 
vector<int> mul(vector<int> a, vector<int> b) {
	int n = 1;
	while (n < a.size() + b.size()) n <<= 1;
	vector<cd> fa(a.begin(), a.end()), fb(b.begin(), b.end());
	fa.resize(n);
	fb.resize(n);
	fft(fa, 0);
	fft(fb, 0);
	for (int i = 0; i < n; i++) fa[i] *= fb[i];
	fft(fa, 1);
	vector<int> r(n);
	for (int i = 0; i < n; i++) r[i] = round(fa[i].real());
	return r;
}
// EndCodeSnip
 
int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);
 
	string s;
	cin >> s;
	int n = s.size();
 
	vector<int> a(n), b(n);
 
	for (int i = 0; i < n; ++i)
		if (s[i] == 'A') a[i] = 1;
		else b[n - i - 1] = 1;
 
	auto c = mul(a, b);
	for (int i = 1; i < n; i++) cout << c[i + n - 1] << '\n';
}
Fundamentos