Optimal Partition (Codeforces 1667B)
Tags: Others, dp

You are given an array $$$a$$$ consisting of $$$n$$$ integers. You should divide $$$a$$$ into continuous non-empty subarrays (there are $$$2^{n-1}$$$ ways to do that).Let $$$s=a_l+a_{l+1}+\ldots+a_r$$$. The value of a subarray $$$a_l, a_{l+1}, \ldots, a_r$$$ is:   $$$(r-l+1)$$$ if $$$s>0$$$, $$$0$$$ if $$$s=0$$$, $$$-(r-l+1)$$$ if $$$s<0$$$.  What is the maximum sum of values you can get with a partition?

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

Output: For each test case print a single integer — the maximum sum of values you can get with an optimal parition.

Note: Test case $$$1$$$: one optimal partition is $$$[1, 2]$$$, $$$[-3]$$$. $$$1+2>0$$$ so the value of $$$[1, 2]$$$ is $$$2$$$. $$$-3<0$$$, so the value of $$$[-3]$$$ is $$$-1$$$. $$$2+(-1)=1$$$.Test case $$$2$$$: the optimal partition is $$$[0, -2, 3]$$$, $$$[-4]$$$, and the sum of values is $$$3+(-1)=2$$$.

Sample Inputs:
5
3
1 2 -3
4
0 -2 3 -4
5
-1 -2 3 -1 -1
6
-1 2 -3 4 -5 6
7
1 -1 -1 1 -1 -1 1

Sample Outputs:
1
2
1
6
-1

Solution (Python):
import os
import sys
from io import BytesIO, IOBase
_str = str
str = lambda x=b'': x if type(x) is bytes else _str(x).encode()
BUFSIZE = 8192

class FastIO(IOBase):
    newlines = 0

    def __init__(self, file):
        self._fd = file.fileno()
        self.buffer = BytesIO()
        self.writable = 'x' in file.mode or 'r' not in file.mode
        self.write = self.buffer.write if self.writable else None

    def read(self):
        while True:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            if not b:
                break
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines = 0
        return self.buffer.read()

    def readline(self):
        while self.newlines == 0:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            self.newlines = b.count(b'\n') + (not b)
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines -= 1
        return self.buffer.readline()

    def flush(self):
        if self.writable:
            os.write(self._fd, self.buffer.getvalue())
            (self.buffer.truncate(0), self.buffer.seek(0))

class IOWrapper(IOBase):

    def __init__(self, file):
        self.buffer = FastIO(file)
        self.flush = self.buffer.flush
        self.writable = self.buffer.writable
        self.write = lambda s: self.buffer.write(s.encode('ascii'))
        self.read = lambda : self.buffer.read().decode('ascii')
        self.readline = lambda : self.buffer.readline().decode('ascii')
(sys.stdin, sys.stdout) = (IOWrapper(sys.stdin), IOWrapper(sys.stdout))
input = lambda : sys.stdin.readline().rstrip('\r\n')
from itertools import accumulate
from bisect import bisect_left
from math import inf

def discrete(a):
    b = sorted(list(set(a)))
    c = [0] * len(a)
    for i in range(len(c)):
        c[i] = bisect_left(b, a[i]) + 1
    return c

def lowbit(x):
    return x & -x

class BIT:

    def __init__(self, x):
        """"""
        self.bit = x
        for i in range(len(x)):
            j = i + lowbit(i)
            if j < len(x):
                x[j] = max(x[j], x[i])

    def update(self, idx, x):
        """"""
        while idx < len(self.bit):
            self.bit[idx] = max(self.bit[idx], x)
            idx += lowbit(idx)

    def query(self, end):
        """"""
        x = -inf
        while end:
            x = max(x, self.bit[end])
            end -= lowbit(end)
        return x

def main():
    t = int(input())
    while t:
        n = int(input())
        a = list(map(int, input().strip().split(' ')))
        pre_sum = [0] + list(accumulate(a))
        pre_sum = discrete(pre_sum)
        bit = BIT([-inf] * (n + 2))
        dp = [0] * (n + 1)
        for i in range(1, n + 1):
            bit.update(pre_sum[i - 1], dp[i - 1] - (i - 1))
            dp[i] = max(dp[i - 1] + (-1 if a[i - 1] < 0 else 0), i + bit.query(pre_sum[i] - 1))
        print(dp[n])
        t -= 1
    return
main()

Solution (Python):
import io
import os
import heapq
from collections import Counter, deque
from array import array
FILE_INPUT_MODE = 0
if FILE_INPUT_MODE >= 0:
    if FILE_INPUT_MODE == 1:
        CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
        FD = os.open(f'{CURRENT_DIR}\\input.txt', os.O_RDONLY)
    elif FILE_INPUT_MODE == 0:
        FD = 0
    DATA = io.BytesIO(os.read(FD, os.fstat(FD).st_size))
    if FD != 0:
        os.close(FD)

    def input():
        return DATA.readline().decode().rstrip('\r\n')
inf = 10 ** 9

class FenwickTree:

    def __init__(self, n):
        self.a = [-inf for _ in range(n + 1)]

    @staticmethod
    def lowbit(x):
        return x & -x

    def update(self, i, x):
        while i < len(self.a):
            self.a[i] = max(self.a[i], x)
            i += self.lowbit(i)

    def query(self, i):
        res = -inf
        while i > 0:
            res = max(res, self.a[i])
            i -= self.lowbit(i)
        return res

def solve(n, a):

    def update(x, i, val):
        nonlocal dp, d, tree_l, tree_r, s
        dp[x] = max(dp[x], val)
        tree_l.update(x, dp[x] - i)
        tree_r.update(len(s) + 1 - x, dp[x] + i)
    pre = [0 for _ in range(n + 1)]
    for i in range(n):
        pre[i + 1] = pre[i] + a[i]
    s = sorted(set(pre[:n + 1]))
    d = {}
    for (i, x) in enumerate(s):
        d[x] = i + 1
    dp = [-inf for _ in range(len(s) + 1)]
    tree_l = FenwickTree(len(s))
    tree_r = FenwickTree(len(s))
    update(d[pre[0]], 0, 0)
    for i in range(1, n + 1):
        x = d[pre[i]]
        mx = dp[x]
        mx = max(mx, tree_l.query(x - 1) + i)
        mx = max(mx, tree_r.query(len(s) - x) - i)
        update(x, i, mx)
    return dp[d[pre[n]]]

def main():
    t = int(input())
    for _ in range(t):
        n = int(input())
        a = [int(word) for word in input().strip().split()]
        print(solve(n, a))
if __name__ == '__main__':
    main()

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def segment_tree(n):
    i = 2
    while True:
        if i >= n * 2:
            tree1 = [-inf] * i
            tree2 = [-inf] * i
            break
        else:
            i *= 2
    return (tree1, tree2)

def update(i, x, tree):
    i += len(tree) // 2
    tree[i] = x
    i //= 2
    while True:
        if i == 0:
            break
        tree[i] = max(tree[2 * i], tree[2 * i + 1])
        i //= 2
    return

def get_max(s, t, tree):
    s += len(tree) // 2
    t += len(tree) // 2
    ans = -inf
    while s <= t:
        if s % 2 == 0:
            s //= 2
        else:
            ans = max(ans, tree[s])
            s = (s + 1) // 2
        if t % 2 == 1:
            t //= 2
        else:
            ans = max(ans, tree[t])
            t = (t - 1) // 2
    return ans
t = int(input())
ans = []
inf = pow(10, 9) + 1
for _ in range(t):
    n = int(input())
    a = list(map(int, input().split()))
    b = [0]
    for i in a:
        b.append(b[-1] + i)
    s = list(set(sorted(b)))
    s.sort()
    d = dict()
    m = len(s)
    for i in range(m):
        d[s[i]] = i + 1
    (tree1, tree2) = segment_tree(m + 5)
    update(d[0], 0, tree1)
    update(d[0], 0, tree2)
    dp = [0] * (n + 1)
    dpmax = [-inf] * (m + 1)
    dpmax[d[0]] = 0
    l = len(tree1) // 2
    for i in range(1, n + 1):
        bi = b[i]
        di = d[bi]
        dpi = max(dpmax[di], get_max(0, di - 1, tree1) + i, get_max(di + 1, m + 2, tree2) - i)
        dpmax[di] = dpi
        update(di, max(tree1[di + l], dpi - i), tree1)
        update(di, max(tree2[di + l], dpi + i), tree2)
        dp[i] = max(dpi, dpmax[di])
    ans0 = dp[n]
    ans.append(ans0)
sys.stdout.write('\n'.join(map(str, ans)))