Sliding Cost

Resumen

Hallar la diferencia entre las sumas de los $K/2$ elementos superiores y los $K/2$ inferiores con dos multiconjuntos.

Solución

El costo de una ventana

Se puede demostrar que es óptimo cambiar los valores de la ventana a la mediana (queda como ejercicio :D). Una vez hallada la mediana, hay que hallar la suma de las distancias de todos los elementos a la mediana. Sumar las distancias de cada elemento por separado es demasiado lento. En su lugar, partiremos los elementos de la ventana en dos grupos y calcularemos el costo como se describe a continuación. Los $K/2$ elementos más pequeños de la ventana estarán en el grupo inferior, mientras que los $K/2$ más grandes estarán en el grupo superior.

El costo de la ventana se puede expresar como una función de $K,S_1,S_2$ y $M$, donde $S_1$ y $S_2$ denotan la suma de los elementos del grupo inferior y del superior respectivamente, y $M$ denota la mediana de la ventana. El costo del grupo inferior será $\sum_{i=1}^{K/2} M-e_i$, y el costo del grupo superior será $\sum_{i=1}^{K/2} e_i-M$, donde $e$ representa un elemento del grupo. Estas expresiones se pueden simplificar a $M\times K/2 - S_1$ y $S_2 - M\times K/2$. El costo total de la ventana es la suma de los costos aportados por ambos grupos, o $S_2-S_1$.

Implementación

Hallar la diferencia entre los $K/2$ elementos más grandes de la ventana y los $K/2$ más pequeños es similar a hallar la mediana deslizante (más información aquí). Para mantener el costo actual, llevamos la cuenta de la suma de cada multiconjunto a medida que insertamos y borramos. Usando el método de doble multiconjunto descrito en la solución de Sliding Median, hacemos que el grupo inferior incluya los $\lceil K/2 \rceil$ elementos inferiores de la ventana. Como resultado, cuando el tamaño de la ventana es impar, el grupo inferior tiene un elemento de más respecto de la cantidad deseada. Podemos corregir esto sumando la mediana a la respuesta final si el tamaño de la ventana es impar.

Complejidad temporal: $\mathcal{O}(N\log K)$

#include <algorithm>
#include <iostream>
#include <set>
using namespace std;
using ll = long long;
const ll mn = (ll)2e5 + 5;
 
ll N, K;
ll arr[mn];
multiset<ll> up;
multiset<ll> low;
ll sLow, sUp;
 
void ins(ll val) {
	ll a = *low.rbegin();
	if (a < val) {
		up.insert(val);
		sUp += val;
		if (up.size() > K / 2) {
			ll moving = *up.begin();
			low.insert(moving);
			sLow += moving;
			up.erase(up.find(moving));
			sUp -= moving;
		}
	} else {
		low.insert(val);
		sLow += val;
		if (low.size() > (K + 1) / 2) {
			ll moving = *low.rbegin();
			up.insert(*low.rbegin());
			sUp += moving;
			low.erase(low.find(*low.rbegin()));
			sLow -= moving;
		}
	}
}
 
void er(ll val) {
	if (up.find(val) != up.end()) up.erase(up.find(val)), sUp -= val;
	else low.erase(low.find(val)), sLow -= val;
	if (low.empty()) {
		ll moving = *up.begin();
		low.insert(*up.begin());
		sLow += moving;
		up.erase(up.find(*up.begin()));
		sUp -= moving;
	}
}
 
ll med() { return (K % 2 == 0) ? 0 : (*low.rbegin()); }
 
int main() {
	cin >> N >> K;
	for (ll i = 0; i < N; i++) cin >> arr[i];
	low.insert(arr[0]);
	sLow += arr[0];
	for (ll i = 1; i < K; i++) ins(arr[i]);
	cout << sUp - sLow + med();
	if (N != 1) cout << " ";
	for (ll i = K; i < N; i++) {
		if (K == 1) {
			ins(arr[i]);
			er(arr[i - K]);
		} else {
			er(arr[i - K]);
			ins(arr[i]);
		}
		cout << sUp - sLow + med();
		if (i != N - 1) cout << " ";
	}
	cout << endl;
}