Autor del problema: | Hasan Saleh (MathModel) |
Preparado por: | Hasan Saleh (MathModel) |
Editorial: | Hasan Saleh (MathModel) |
Una solución simple es iterar a través de todos los subsegmentos en la matriz y sumar las puntuaciones, lo que da una solución de \mathcal{O}(n^2). Sin embargo, esto es demasiado lento; necesitamos hacerlo más rápido.
Para acelerar el cálculo, necesitamos calcular para cada subsegmento que termina en el índice , cuántos de ellos son maravillosos, cuántos son perfectos y cuántos son normales.
Sea la suma de la matriz. Podemos determinar fácilmente si la suma del subsegmento es mayor que \frac{s}{2}, será maravillosa; si es menor que \frac{s}{2}, será normal, y si es par y la suma del subsegmento es igual a \frac{s}{2}, será perfecta.
Así, necesitamos obtener la suma prefija para cada índice p_i, y contar el número de segmentos que son mayores, iguales o menores que p_i - \frac{s}{2}. Esto se puede hacer de muchas maneras; una forma es calcular todas las sumas prefijas y mapearlas a índices de 1-m, donde es el número de sumas prefijas distintas. Luego podemos usar un árbol de segmentos o un árbol de Fenwick junto con búsqueda binaria para resolver este problema. Esto hará que la complejidad temporal sea \mathcal{O}(n \log n), lo que es lo suficientemente rápido para resolver esto usando Python.
import bisect
class fenwick(object):
def __init__(self, n):
self.n = n
self.cul = [0]*n
self.maxd = len(bin(n))-3
def update(self,index,diff):
i = index
while i<self.n:
self.cul[i] += diff
i += (i+1)&(-i-1)
def getaccu(self,index):
output = 0
i = index
while i>=0:
output += self.cul[i]
i -= (i+1)&(-i-1)
return output
def solve():
n = int(input())
arr = list(map(int,input().split()))
tot = sum(arr)
nums = [0]
for num in arr:
nums.append(nums[-1]+num)
nums = sorted(list(set(nums)))
nums = [-float('inf')] + nums
mapping = {}
already = {}
already[0] = 1
for i in range(len(nums)):
mapping[nums[i]] = i
fen = fenwick(len(nums)+1)
fen.update(mapping[0],1)
count = 0
accu = 0
ans = 0
for i,num in enumerate(arr):
accu += num
target = accu + (-tot) // 2
loc = bisect.bisect(nums, target) - 1
good = fen.getaccu(loc)
ans += good * 2
ans -= i + 1 - good
if tot % 2 == 0 and target in already:
ans += already[target]
fen.update(mapping[accu],1)
if accu not in already: already[accu] = 0
already[accu] += 1
print(ans)
for _ in range(int(input())):
solve()