Score#11806

Score (Corrigé)

Auteur du problème:	Hasan Saleh (MathModel)
Préparé par:	Hasan Saleh (MathModel)
Éditorial:	Hasan Saleh (MathModel)

Une solution simple consiste à parcourir tous les sous-segments du tableau et à additionner les scores, ce qui donne une solution en $O (n^{2})$ \mathcal{O}(n^2). Cependant, cela est trop lent ; nous devons le faire plus rapidement.

Pour accélérer le calcul, nous devons calculer pour chaque sous-segment qui se termine à l'index $i$ , combien d'entre eux sont merveilleux, combien sont parfaits, et combien sont normaux.

Soit $s$ la somme du tableau. Nous pouvons facilement déterminer si la somme du sous-segment est supérieure à $\frac{s}{2}$ \frac{s}{2}, elle sera merveilleuse ; si elle est inférieure à $\frac{s}{2}$ \frac{s}{2}, elle sera normale, et si $s$ est pair et que la somme du sous-segment est égale à $\frac{s}{2}$ \frac{s}{2}, elle sera parfaite.

Ainsi, nous devons obtenir la somme préfixe pour chaque index $p_{i}$ p_i, et compter le nombre de segments qui sont supérieurs, égaux ou inférieurs à $p_{i} - \frac{s}{2}$ p_i - \frac{s}{2}. Cela peut être fait de plusieurs manières ; une façon consiste à calculer toutes les sommes préfixes et à les mapper à des indices de $1 - m$ 1-m, où $m$ est le nombre de sommes préfixes distinctes. Ensuite, nous pouvons utiliser un arbre de segments ou un arbre de Fenwick avec une recherche binaire pour résoudre ce problème. Cela rendra la complexité temporelle $O (n lo g n)$ \mathcal{O}(n \log n), ce qui est suffisamment rapide pour résoudre cela en 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()

Python