Zero-Sum Prefixes (Codeforces 1748C)
Tags: Greedy, data structures

The score of an array $$$v_1,v_2,\ldots,v_n$$$ is defined as the number of indices $$$i$$$ ($$$1 \le i \le n$$$) such that $$$v_1+v_2+\ldots+v_i = 0$$$.You are given an array $$$a_1,a_2,\ldots,a_n$$$ of length $$$n$$$. You can perform the following operation multiple times:  select an index $$$i$$$ ($$$1 \le i \le n$$$) such that $$$a_i=0$$$;  then replace $$$a_i$$$ by an arbitrary integer.  What is the maximum possible score of $$$a$$$ that can be obtained by performing a sequence of such operations?

Input: Each test contains multiple test cases. The first line contains a single integer $$$t$$$ ($$$1 \le t \le 10^4$$$) — the number of test cases.The first line of each test case contains one integer $$$n$$$ ($$$1 \le n \le 2 \cdot 10^5$$$) — the length of the array $$$a$$$.The second line of each test case contains $$$n$$$ integers $$$a_1,a_2,\ldots,a_n$$$ ($$$-10^9 \le a_i \le 10^9$$$) — array $$$a$$$.It is guaranteed that the sum of $$$n$$$ over all test cases does not exceed $$$2 \cdot 10^5$$$.

Output: For each test case, print the maximum possible score of the array $$$a$$$ after performing a sequence of operations.

Note: In the first test case, it is optimal to change the value of $$$a_2$$$ to $$$-2$$$ in one operation.The resulting array $$$a$$$ will be $$$[2,-2,1,-1,0]$$$, with a score of $$$3$$$:  $$$a_1+a_2=2-2=0$$$;  $$$a_1+a_2+a_3+a_4=2-2+1-1=0$$$;  $$$a_1+a_2+a_3+a_4+a_5=2-2+1-1+0=0$$$. In the second test case, it is optimal to change the value of $$$a_3$$$ to $$$-2\,000\,000\,000$$$, giving us an array with a score of $$$1$$$.In the third test case, it is not necessary to perform any operations.

Sample Inputs:
5
5
2 0 1 -1 0
3
1000000000 1000000000 0
4
0 0 0 0
8
3 0 2 -10 10 -30 30 0
9
1 0 0 1 -1 0 1 0 -1

Sample Outputs:
3
1
4
4
5

Solution (Python):
import sys
from itertools import accumulate
input = sys.stdin.readline
from collections import defaultdict

def solve():
    n = int(input())
    a = list(map(int, input().split()))
    preSum = list(accumulate(a))
    ans = 0
    d = defaultdict(int)
    maxn = 0
    for i in range(n - 1, -1, -1):
        d[preSum[i]] += 1
        maxn = max(maxn, d[preSum[i]])
        if a[i] == 0:
            ans += maxn
            d = defaultdict(int)
            maxn = 0
    ans += d[0]
    return print(ans)
if __name__ == '__main__':
    for _ in range(int(input())):
        solve()

Solution (Python):
import os
import heapq
import collections
import sys
import math
import bisect
from io import BytesIO, IOBase
input = sys.stdin.readline

def solve():
    n = int(input())
    a = list(map(int, input().split()))
    pre = [a[0]] * n
    d = collections.defaultdict(lambda : 0)
    ans = 0
    for i in range(1, n):
        pre[i] = pre[i - 1] + a[i]
    for i in reversed(range(n)):
        d[pre[i]] += 1
        if a[i] == 0:
            ans += max(d.values())
            d.clear()
    ans += d[0]
    print(ans)
if __name__ == '__main__':
    testcases = 1
    if testcases:
        for i in range(int(input())):
            solve()
    else:
        solve()

Solution (Python):
import collections
import heapq
import sys
import math
import bisect
import os
from io import BytesIO, IOBase
input = sys.stdin.readline

def solve():
    n = int(input())
    l = list(map(int, input().split()))
    dp = [0] * n
    dp[0] = l[0]
    d = collections.defaultdict(int)
    tot = 0
    for i in range(1, n):
        dp[i] = dp[i - 1] + l[i]
    for j in reversed(range(n)):
        d[dp[j]] += 1
        if l[j] == 0:
            tot += max(d.values())
            d.clear()
    tot += d[0]
    print(tot)
if __name__ == '__main__':
    for i in range(int(input())):
        solve()