Balance Addicts (Codeforces 1738E)
Tags: Math, two pointers

Given an integer sequence $$$a_1, a_2, \dots, a_n$$$ of length $$$n$$$, your task is to compute the number, modulo $$$998244353$$$, of ways to partition it into several non-empty continuous subsequences such that the sums of elements in the subsequences form a balanced sequence.A sequence $$$s_1, s_2, \dots, s_k$$$ of length $$$k$$$ is said to be balanced, if $$$s_{i} = s_{k-i+1}$$$ for every $$$1 \leq i \leq k$$$. For example, $$$[1, 2, 3, 2, 1]$$$ and $$$[1,3,3,1]$$$ are balanced, but $$$[1,5,15]$$$ is not. Formally, every partition can be described by a sequence of indexes $$$i_1, i_2, \dots, i_k$$$ of length $$$k$$$ with $$$1 = i_1 < i_2 < \dots < i_k \leq n$$$ such that   $$$k$$$ is the number of non-empty continuous subsequences in the partition;  For every $$$1 \leq j \leq k$$$, the $$$j$$$-th continuous subsequence starts with $$$a_{i_j}$$$, and ends exactly before $$$a_{i_{j+1}}$$$, where $$$i_{k+1} = n + 1$$$. That is, the $$$j$$$-th subsequence is $$$a_{i_j}, a_{i_j+1}, \dots, a_{i_{j+1}-1}$$$.  There are $$$2^{n-1}$$$ different partitions in total. Let $$$s_1, s_2, \dots, s_k$$$ denote the sums of elements in the subsequences with respect to the partition $$$i_1, i_2, \dots, i_k$$$. Formally, for every $$$1 \leq j \leq k$$$, $$$$$$ s_j = \sum_{i=i_{j}}^{i_{j+1}-1} a_i = a_{i_j} + a_{i_j+1} + \dots + a_{i_{j+1}-1}. $$$$$$ For example, the partition $$$[1\,|\,2,3\,|\,4,5,6]$$$ of sequence $$$[1,2,3,4,5,6]$$$ is described by the sequence $$$[1,2,4]$$$ of indexes, and the sums of elements in the subsequences with respect to the partition is $$$[1,5,15]$$$.Two partitions $$$i_1, i_2, \dots, i_k$$$ and $$$i'_1, i'_2, \dots, i'_{k'}$$$ (described by sequences of indexes) are considered to be different, if at least one of the following holds.   $$$k \neq k'$$$,  $$$i_j \neq i'_j$$$ for some $$$1 \leq j \leq \min\left\{ k, k' \right\}$$$. 

Input: Each test contains multiple test cases. The first line contains an integer $$$t$$$ ($$$1 \leq t \leq 10^5$$$) — the number of test cases. The following lines contain the description of each test case.The first line of each test case contains an integer $$$n$$$ ($$$1 \leq n \leq 10^5$$$), indicating the length of the sequence $$$a$$$.The second line of each test case contains $$$n$$$ integers $$$a_1, a_2, \dots, a_n$$$ ($$$0 \leq a_i \leq 10^9$$$), indicating the elements of the sequence $$$a$$$.It is guaranteed that the sum of $$$n$$$ over all test cases does not exceed $$$10^5$$$.

Output: For each test case, output the number of partitions with respect to which the sum of elements in each subsequence is balanced, modulo $$$998244353$$$.

Note: For the first test case, there is only one way to partition a sequence of length $$$1$$$, which is itself and is, of course, balanced. For the second test case, there are $$$2$$$ ways to partition it:   The sequence $$$[1, 1]$$$ itself, then $$$s = [2]$$$ is balanced;  Partition into two subsequences $$$[1\,|\,1]$$$, then $$$s = [1, 1]$$$ is balanced. For the third test case, there are $$$3$$$ ways to partition it:   The sequence $$$[0, 0, 1, 0]$$$ itself, then $$$s = [1]$$$ is balanced;  $$$[0 \,|\, 0, 1 \,|\, 0]$$$, then $$$s = [0, 1, 0]$$$ is balanced;  $$$[0, 0 \,|\, 1 \,|\, 0]$$$, then $$$s = [0, 1, 0]$$$ is balanced. For the fourth test case, there are $$$4$$$ ways to partition it:   The sequence $$$[1, 2, 3, 2, 1]$$$ itself, then $$$s = [9]$$$ is balanced;  $$$[1, 2 \,|\, 3 \,|\, 2, 1]$$$, then $$$s = [3, 3, 3]$$$ is balanced;  $$$[1 \,|\, 2, 3, 2 \,|\, 1]$$$, then $$$s = [1, 7, 1]$$$ is balanced;  $$$[1 \,|\, 2 \,|\, 3 \,|\, 2 \,|\, 1]$$$, then $$$s = [1, 2, 3, 2, 1]$$$ is balanced. For the fifth test case, there are $$$2$$$ ways to partition it:   The sequence $$$[1, 3, 5, 7, 9]$$$ itself, then $$$s = [25]$$$ is balanced;  $$$[1, 3, 5 \,|\, 7 \,|\, 9]$$$, then $$$s = [9, 7, 9]$$$ is balanced. For the sixth test case, every possible partition should be counted. So the answer is $$$2^{32-1} \equiv 150994942 \pmod {998244353}$$$.

Sample Inputs:
6
1
1000000000
2
1 1
4
0 0 1 0
5
1 2 3 2 1
5
1 3 5 7 9
32
0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0

Sample Outputs:
1
2
3
4
2
150994942

Solution (Python):
MOD = 998244353
MAX = 100005
fact = [1]
for zz in range(1, MAX + 1):
    fact.append(fact[-1] * zz % MOD)

def nCr(n, r):
    num = fact[n]
    den = fact[r] * fact[n - r] % MOD
    return num * pow(den, MOD - 2, MOD) % MOD

def main():
    t = int(input())
    allans = []
    for _ in range(t):
        n = int(input())
        a = readIntArr()
        p = a.copy()
        for i in range(1, n):
            p[i] += p[i - 1]

        def q(l, r):
            if r + 1 == l:
                return 0
            if l == 0:
                return p[r]
            return p[r] - p[l - 1]
        cnts = 1
        l = 0
        r = n - 1
        while l + 1 <= r:
            while q(0, l) > q(r, n - 1):
                r -= 1
                if l + 1 > r:
                    break
            if l + 1 > r:
                break
            if q(0, l) == q(r, n - 1):
                if q(l + 1, r - 1) == 0:
                    gaps = r - l
                    cnts2 = 0
                    for r in range(1, gaps + 1):
                        cnts2 += nCr(gaps, r)
                        cnts2 %= MOD
                    cnts = cnts + cnts * cnts2 % MOD
                    cnts %= MOD
                    break
                l2 = l
                left_cnt = 1
                while a[l2 + 1] == 0:
                    left_cnt += 1
                    l2 += 1
                r2 = r
                right_cnt = 1
                while a[r2 - 1] == 0:
                    right_cnt += 1
                    r2 -= 1
                cnts2 = 0
                z = min(left_cnt, right_cnt)
                for y in range(1, z + 1):
                    cnts2 += nCr(left_cnt, y) * nCr(right_cnt, y) % MOD
                    cnts2 %= MOD
                cnts = cnts + cnts * cnts2 % MOD
                cnts %= MOD
                l = l2
                r = r2
            l += 1
        allans.append(cnts)
    multiLineArrayPrint(allans)
    return
import sys
input = sys.stdin.buffer.readline

def oneLineArrayPrint(arr):
    print(' '.join([str(x) for x in arr]))

def multiLineArrayPrint(arr):
    print('\n'.join([str(x) for x in arr]))

def multiLineArrayOfArraysPrint(arr):
    print('\n'.join([' '.join([str(x) for x in y]) for y in arr]))

def readIntArr():
    return [int(x) for x in input().split()]

def makeArr(defaultValFactory, dimensionArr):
    dv = defaultValFactory
    da = dimensionArr
    if len(da) == 1:
        return [dv() for _ in range(da[0])]
    else:
        return [makeArr(dv, da[1:]) for _ in range(da[0])]

def queryInteractive(a, b):
    print('? {} {}'.format(a, b))
    sys.stdout.flush()
    return int(input())

def answerInteractive(ans):
    print('! {}'.format(ans))
    sys.stdout.flush()
inf = float('inf')
from math import gcd, floor, ceil
import math
for _abc in range(1):
    main()

Solution (Python):
import io, os
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline
M = 998244353
MAXNUM = 100020
factor = [1] * (MAXNUM + 2)
factor[-1] = 0
storecomb = {}
for i in range(1, MAXNUM + 1):
    factor[i] = factor[i - 1] * i % M

def fastfrac(a, b):
    numb = pow(b, M - 2, M)
    return a % M * (numb % M) % M

def comb(n, k):
    if n < k:
        return 0
    if n == k:
        return 1
    if (n, k) in storecomb:
        return storecomb[n, k]
    num1 = factor[n]
    num2 = factor[k]
    num3 = factor[n - k]
    num = num2 * num3 % M
    output = fastfrac(num1, num)
    storecomb[n, k] = output
    return output

def main(t):
    n = int(input())
    arr = list(map(int, input().split()))
    front = 0
    rear = n - 1
    pick = [[-1]]
    temp = []
    frontaccu = 0
    rearaccu = 0
    while front < n:
        while front < n and arr[front] == 0:
            front += 1
        while rear >= 0 and arr[rear] == 0:
            rear -= 1
        if front >= n:
            break
        if frontaccu < rearaccu:
            temp.append(front)
            frontaccu += arr[front]
            front += 1
        elif frontaccu > rearaccu:
            rearaccu += arr[rear]
            rear -= 1
        else:
            if temp:
                pick.append(temp)
            temp = []
            frontaccu = arr[front]
            rearaccu = arr[rear]
            temp.append(front)
            front += 1
            rear -= 1
    if temp:
        pick.append(temp)
    pick.append([n])
    nums = []
    for i in range(len(pick) - 1):
        nums.append(pick[i + 1][0] - pick[i][-1])
    nums[0] -= 1
    nums[-1] -= 1
    ans = 1
    front = 0
    rear = len(nums) - 1
    while front <= rear:
        bei = 0
        if front < rear:
            minimum = min(nums[front], nums[rear])
            for i in range(minimum + 1):
                bei += comb(nums[front], i) * comb(nums[rear], i)
                bei = bei % M
        else:
            bei = 1 << nums[front]
        ans = ans * bei % M
        front += 1
        rear -= 1
    print(ans % M)
T = int(input())
t = 1
while t <= T:
    main(t)
    t += 1

Solution (Python):
import sys
input = sys.stdin.buffer.readline
p = 998244353
fact_dict = [1, 1]
for i in range(2, 10 ** 5 + 1):
    fact_dict.append(fact_dict[-1] * i % p)

def choose(a, b):
    if b > a:
        return 0
    num = fact_dict[a]
    den = fact_dict[b] * fact_dict[a - b] % p
    den1 = pow(den, p - 2, p)
    return num * den1 % p

def process1(L):
    counts = [0, 0]
    ''
    for (i, t) in L:
        counts[t] += 1
    curr = 0
    pref = 0
    suff = counts[0]
    answer = 0
    for (i, t) in L:
        if t == 1:
            pref += 1
        else:
            ''
            if pref > 0:
                entry = choose(pref + suff - 1, pref - 1)
                answer = (answer + entry) % p
            suff -= 1
    return answer

def process(A):
    n = len(A)
    pref = [0 for i in range(n + 1)]
    suff = [0 for i in range(n + 1)]
    for i in range(n):
        pref[i + 1] = pref[i] + A[i]
    for i in range(n - 1, -1, -1):
        suff[i] = suff[i + 1] + A[i]
    L = []
    for i in range(1, n + 1):
        L.append([pref[i], i - 1, 1])
        ''
    for i in range(n):
        L.append([suff[i], i, 0])
        ''
    L.sort()
    answer = 1
    curr = [L[0][0], []]
    for i in range(len(L)):
        (x, I, t) = L[i]
        if x > curr[0]:
            entry = process1(curr[1])
            answer = answer * (1 + entry) % p
            curr = [x, []]
        if x == curr[0]:
            curr[1].append([I, t])
    entry = process1(curr[1])
    answer = answer * (1 + entry) % p
    sys.stdout.write(f'{answer}\n')
    return
t = int(input())
for i in range(t):
    n = int(input())
    A = [int(x) for x in input().split()]
    process(A)