Optimal Partition (Codeforces 1668D)
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 math
import random
import heapq, bisect
import sys
from collections import deque, defaultdict
from fractions import Fraction
from itertools import permutations
from collections import defaultdict
mod = 10 ** 9 + 7
mod1 = 998244353
import sys
import os
import sys
from io import BytesIO, IOBase
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')

class SegmentTree:

    def __init__(self, data, default=-10 ** 9, func=lambda a, b: max(a, b)):
        """"""
        self._default = default
        self._func = func
        self._len = len(data)
        self._size = _size = 1 << (self._len - 1).bit_length()
        self.data = [default] * (2 * _size)
        self.data[_size:_size + self._len] = data
        for i in reversed(range(_size)):
            self.data[i] = func(self.data[i + i], self.data[i + i + 1])

    def __delitem__(self, idx):
        self[idx] = self._default

    def __getitem__(self, idx):
        return self.data[idx + self._size]

    def __setitem__(self, idx, value):
        idx += self._size
        self.data[idx] = value
        idx >>= 1
        while idx:
            self.data[idx] = self._func(self.data[2 * idx], self.data[2 * idx + 1])
            idx >>= 1

    def __len__(self):
        return self._len

    def query(self, start, stop):
        if start == stop:
            return self.__getitem__(start)
        stop += 1
        start += self._size
        stop += self._size
        res = self._default
        while start < stop:
            if start & 1:
                res = self._func(res, self.data[start])
                start += 1
            if stop & 1:
                stop -= 1
                res = self._func(res, self.data[stop])
            start >>= 1
            stop >>= 1
        return res

    def __repr__(self):
        return 'SegmentTree({0})'.format(self.data)

class SortedList:

    def __init__(self, iterable=[], _load=200):
        """"""
        values = sorted(iterable)
        self._len = _len = len(values)
        self._load = _load
        self._lists = _lists = [values[i:i + _load] for i in range(0, _len, _load)]
        self._list_lens = [len(_list) for _list in _lists]
        self._mins = [_list[0] for _list in _lists]
        self._fen_tree = []
        self._rebuild = True

    def _fen_build(self):
        """"""
        self._fen_tree[:] = self._list_lens
        _fen_tree = self._fen_tree
        for i in range(len(_fen_tree)):
            if i | i + 1 < len(_fen_tree):
                _fen_tree[i | i + 1] += _fen_tree[i]
        self._rebuild = False

    def _fen_update(self, index, value):
        """"""
        if not self._rebuild:
            _fen_tree = self._fen_tree
            while index < len(_fen_tree):
                _fen_tree[index] += value
                index |= index + 1

    def _fen_query(self, end):
        """"""
        if self._rebuild:
            self._fen_build()
        _fen_tree = self._fen_tree
        x = 0
        while end:
            x += _fen_tree[end - 1]
            end &= end - 1
        return x

    def _fen_findkth(self, k):
        """"""
        _list_lens = self._list_lens
        if k < _list_lens[0]:
            return (0, k)
        if k >= self._len - _list_lens[-1]:
            return (len(_list_lens) - 1, k + _list_lens[-1] - self._len)
        if self._rebuild:
            self._fen_build()
        _fen_tree = self._fen_tree
        idx = -1
        for d in reversed(range(len(_fen_tree).bit_length())):
            right_idx = idx + (1 << d)
            if right_idx < len(_fen_tree) and k >= _fen_tree[right_idx]:
                idx = right_idx
                k -= _fen_tree[idx]
        return (idx + 1, k)

    def _delete(self, pos, idx):
        """"""
        _lists = self._lists
        _mins = self._mins
        _list_lens = self._list_lens
        self._len -= 1
        self._fen_update(pos, -1)
        del _lists[pos][idx]
        _list_lens[pos] -= 1
        if _list_lens[pos]:
            _mins[pos] = _lists[pos][0]
        else:
            del _lists[pos]
            del _list_lens[pos]
            del _mins[pos]
            self._rebuild = True

    def _loc_left(self, value):
        """"""
        if not self._len:
            return (0, 0)
        _lists = self._lists
        _mins = self._mins
        (lo, pos) = (-1, len(_lists) - 1)
        while lo + 1 < pos:
            mi = lo + pos >> 1
            if value <= _mins[mi]:
                pos = mi
            else:
                lo = mi
        if pos and value <= _lists[pos - 1][-1]:
            pos -= 1
        _list = _lists[pos]
        (lo, idx) = (-1, len(_list))
        while lo + 1 < idx:
            mi = lo + idx >> 1
            if value <= _list[mi]:
                idx = mi
            else:
                lo = mi
        return (pos, idx)

    def _loc_right(self, value):
        """"""
        if not self._len:
            return (0, 0)
        _lists = self._lists
        _mins = self._mins
        (pos, hi) = (0, len(_lists))
        while pos + 1 < hi:
            mi = pos + hi >> 1
            if value < _mins[mi]:
                hi = mi
            else:
                pos = mi
        _list = _lists[pos]
        (lo, idx) = (-1, len(_list))
        while lo + 1 < idx:
            mi = lo + idx >> 1
            if value < _list[mi]:
                idx = mi
            else:
                lo = mi
        return (pos, idx)

    def add(self, value):
        """"""
        _load = self._load
        _lists = self._lists
        _mins = self._mins
        _list_lens = self._list_lens
        self._len += 1
        if _lists:
            (pos, idx) = self._loc_right(value)
            self._fen_update(pos, 1)
            _list = _lists[pos]
            _list.insert(idx, value)
            _list_lens[pos] += 1
            _mins[pos] = _list[0]
            if _load + _load < len(_list):
                _lists.insert(pos + 1, _list[_load:])
                _list_lens.insert(pos + 1, len(_list) - _load)
                _mins.insert(pos + 1, _list[_load])
                _list_lens[pos] = _load
                del _list[_load:]
                self._rebuild = True
        else:
            _lists.append([value])
            _mins.append(value)
            _list_lens.append(1)
            self._rebuild = True

    def discard(self, value):
        """"""
        _lists = self._lists
        if _lists:
            (pos, idx) = self._loc_right(value)
            if idx and _lists[pos][idx - 1] == value:
                self._delete(pos, idx - 1)

    def remove(self, value):
        """"""
        _len = self._len
        self.discard(value)
        if _len == self._len:
            raise ValueError('{0!r} not in list'.format(value))

    def pop(self, index=-1):
        """"""
        (pos, idx) = self._fen_findkth(self._len + index if index < 0 else index)
        value = self._lists[pos][idx]
        self._delete(pos, idx)
        return value

    def bisect_left(self, value):
        """"""
        (pos, idx) = self._loc_left(value)
        return self._fen_query(pos) + idx

    def bisect_right(self, value):
        """"""
        (pos, idx) = self._loc_right(value)
        return self._fen_query(pos) + idx

    def count(self, value):
        """"""
        return self.bisect_right(value) - self.bisect_left(value)

    def __len__(self):
        """"""
        return self._len

    def __getitem__(self, index):
        """"""
        (pos, idx) = self._fen_findkth(self._len + index if index < 0 else index)
        return self._lists[pos][idx]

    def __delitem__(self, index):
        """"""
        (pos, idx) = self._fen_findkth(self._len + index if index < 0 else index)
        self._delete(pos, idx)

    def __contains__(self, value):
        """"""
        _lists = self._lists
        if _lists:
            (pos, idx) = self._loc_left(value)
            return idx < len(_lists[pos]) and _lists[pos][idx] == value
        return False

    def __iter__(self):
        """"""
        return (value for _list in self._lists for value in _list)

    def __reversed__(self):
        """"""
        return (value for _list in reversed(self._lists) for value in reversed(_list))

    def __repr__(self):
        """"""
        return 'SortedList({0})'.format(list(self))

class Factorial:

    def __init__(self, MOD):
        self.MOD = MOD
        self.factorials = [1, 1]
        self.invModulos = [0, 1]
        self.invFactorial_ = [1, 1]

    def calc(self, n):
        if n <= -1:
            print('Invalid argument to calculate n!')
            print('n must be non-negative value. But the argument was ' + str(n))
            exit()
        if n < len(self.factorials):
            return self.factorials[n]
        nextArr = [0] * (n + 1 - len(self.factorials))
        initialI = len(self.factorials)
        prev = self.factorials[-1]
        m = self.MOD
        for i in range(initialI, n + 1):
            prev = nextArr[i - initialI] = prev * i % m
        self.factorials += nextArr
        return self.factorials[n]

    def inv(self, n):
        if n <= -1:
            print('Invalid argument to calculate n^(-1)')
            print('n must be non-negative value. But the argument was ' + str(n))
            exit()
        p = self.MOD
        pi = n % p
        if pi < len(self.invModulos):
            return self.invModulos[pi]
        nextArr = [0] * (n + 1 - len(self.invModulos))
        initialI = len(self.invModulos)
        for i in range(initialI, min(p, n + 1)):
            next = -self.invModulos[p % i] * (p // i) % p
            self.invModulos.append(next)
        return self.invModulos[pi]

    def invFactorial(self, n):
        if n <= -1:
            print('Invalid argument to calculate (n^(-1))!')
            print('n must be non-negative value. But the argument was ' + str(n))
            exit()
        if n < len(self.invFactorial_):
            return self.invFactorial_[n]
        self.inv(n)
        nextArr = [0] * (n + 1 - len(self.invFactorial_))
        initialI = len(self.invFactorial_)
        prev = self.invFactorial_[-1]
        p = self.MOD
        for i in range(initialI, n + 1):
            prev = nextArr[i - initialI] = prev * self.invModulos[i % p] % p
        self.invFactorial_ += nextArr
        return self.invFactorial_[n]

class Combination:

    def __init__(self, MOD):
        self.MOD = MOD
        self.factorial = Factorial(MOD)

    def ncr(self, n, k):
        if k < 0 or n < k:
            return 0
        k = min(k, n - k)
        f = self.factorial
        return f.calc(n) * f.invFactorial(max(n - k, k)) * f.invFactorial(min(k, n - k)) % self.MOD

def powm(a, n, m):
    if a == 1 or n == 0:
        return 1
    if n % 2 == 0:
        s = powm(a, n // 2, m)
        return s * s % m
    else:
        return a * powm(a, n - 1, m) % m

def sort_list(list1, list2):
    zipped_pairs = zip(list2, list1)
    z = [x for (_, x) in sorted(zipped_pairs)]
    return z

def product(l):
    por = 1
    for i in range(len(l)):
        por *= l[i]
    return por

def binarySearchCount(arr, n, key):
    left = 0
    right = n - 1
    count = 0
    while left <= right:
        mid = int((right + left) / 2)
        if arr[mid] <= key:
            count = mid + 1
            left = mid + 1
        else:
            right = mid - 1
    return count

def countdig(n):
    c = 0
    while n > 0:
        n //= 10
        c += 1
    return c

def binary(x, length):
    y = bin(x)[2:]
    return y if len(y) >= length else '0' * (length - len(y)) + y

def countGreater(arr, n, k):
    l = 0
    r = n - 1
    leftGreater = n
    while l <= r:
        m = int(l + (r - l) / 2)
        if arr[m] >= k:
            leftGreater = m
            r = m - 1
        else:
            l = m + 1
    return n - leftGreater

class TrieNode:

    def __init__(self):
        self.children = [None] * 26
        self.isEndOfWord = False

class Trie:

    def __init__(self):
        self.root = self.getNode()

    def getNode(self):
        return TrieNode()

    def _charToIndex(self, ch):
        return ord(ch) - ord('a')

    def insert(self, key):
        pCrawl = self.root
        length = len(key)
        for level in range(length):
            index = self._charToIndex(key[level])
            if not pCrawl.children[index]:
                pCrawl.children[index] = self.getNode()
            pCrawl = pCrawl.children[index]
        pCrawl.isEndOfWord = True

    def search(self, key):
        pCrawl = self.root
        length = len(key)
        for level in range(length):
            index = self._charToIndex(key[level])
            if not pCrawl.children[index]:
                return False
            pCrawl = pCrawl.children[index]
        return pCrawl != None and pCrawl.isEndOfWord

class Node:

    def __init__(self, data):
        self.data = data
        self.count = 0
        self.left = None
        self.right = None

class BinaryTrie:

    def __init__(self):
        self.root = Node(0)

    def insert(self, pre_xor):
        self.temp = self.root
        for i in range(31, -1, -1):
            val = pre_xor & 1 << i
            if val:
                if not self.temp.right:
                    self.temp.right = Node(0)
                self.temp = self.temp.right
                self.temp.count += 1
            if not val:
                if not self.temp.left:
                    self.temp.left = Node(0)
                self.temp = self.temp.left
                self.temp.count += 1
        self.temp.data = pre_xor

    def query(self, xor):
        self.temp = self.root
        for i in range(31, -1, -1):
            val = xor & 1 << i
            if not val:
                if self.temp.left and self.temp.left.count > 0:
                    self.temp = self.temp.left
                elif self.temp.right:
                    self.temp = self.temp.right
            elif self.temp.right and self.temp.right.count > 0:
                self.temp = self.temp.right
            elif self.temp.left:
                self.temp = self.temp.left
            self.temp.count -= 1
        return xor ^ self.temp.data
for ik in range(int(input())):
    n = int(input())
    l = list(map(int, input().split()))
    pre = [0]
    q = defaultdict(int)
    q[0] = 0
    for i in range(n):
        pre.append(pre[-1] + l[i])
        if pre[-1] not in q:
            q[pre[-1]] = i + 1
    s = set(pre)
    s = sorted(s)
    ind = defaultdict(int)
    for i in range(len(s)):
        ind[s[i]] = i
    dp1 = [-10 ** 9] * (n + 1)
    dp1[ind[0]] = 0
    s = SegmentTree(dp1)
    dp = [-10 ** 9] * (n + 1)
    dp[0] = 0
    for i in range(1, n + 1):
        dp[i] = max(dp[i], s.query(0, ind[pre[i]] - 1) + i)
        if l[i - 1] < 0:
            dp[i] = max(dp[i], dp[i - 1] - 1)
        if l[i - 1] == 0:
            dp[i] = max(dp[i], dp[i - 1])
        dp1[ind[pre[i]]] = max(dp1[ind[pre[i]]], dp[i] - i)
        s.__setitem__(ind[pre[i]], dp1[ind[pre[i]]])
    print(dp[-1])

Solution (Python):
from sys import stdin, stdout
MAXN = int(500000.0 + 5)
INF = int(1e+18)
(c, dp, pre) = [[0 for _ in range(MAXN)] for i in range(3)]
lis = []

def lowbit(x):
    return x & -x

def update(x, k, n):
    while x <= n:
        c[x] = max(c[x], k)
        x += lowbit(x)

def getmax(x):
    ans = -INF
    while x > 0:
        ans = max(ans, c[x])
        x -= lowbit(x)
    return ans
t = int(stdin.readline())
for _ in range(t):
    n = int(stdin.readline())
    id = dict()
    st = set()
    ax = [int(i) for i in stdin.readline().split()]
    a = [0] * (n + 1)
    st.add(0)
    for i in range(1, n + 1):
        c[i] = -INF
        dp[i] = -INF
        a[i] = ax[i - 1]
        pre[i] = pre[i - 1] + a[i]
        st.add(pre[i])
    cnt = 0
    lis = list(st)
    lis.sort()
    for x in lis:
        cnt += 1
        id[x] = cnt
    update(id[0], 0, n)
    for i in range(1, n + 1):
        if a[i] < 0:
            dp[i] = dp[i - 1] - 1
        elif a[i] == 0:
            dp[i] = dp[i - 1]
        dp[i] = max(dp[i], getmax(id[pre[i]] - 1) + i)
        update(id[pre[i]], dp[i] - i, n)
    print(dp[n])

Solution (Python):
from sys import stdin, stdout

class SegmentTree:

    def __init__(self, data, default=-10 ** 6, func=max):
        """"""
        self._default = default
        self._func = func
        self._len = len(data)
        self._size = _size = 1 << (self._len - 1).bit_length()
        self.data = [default] * (2 * _size)
        self.data[_size:_size + self._len] = data
        for i in reversed(range(_size)):
            self.data[i] = func(self.data[i + i], self.data[i + i + 1])

    def __delitem__(self, idx):
        self[idx] = self._default

    def __getitem__(self, idx):
        return self.data[idx + self._size]

    def __setitem__(self, idx, value):
        idx += self._size
        self.data[idx] = value
        idx >>= 1
        while idx:
            self.data[idx] = self._func(self.data[2 * idx], self.data[2 * idx + 1])
            idx >>= 1

    def __len__(self):
        return self._len

    def query(self, start, stop):
        """"""
        start += self._size
        stop += self._size
        res_left = res_right = self._default
        while start < stop:
            if start & 1:
                res_left = self._func(res_left, self.data[start])
                start += 1
            if stop & 1:
                stop -= 1
                res_right = self._func(self.data[stop], res_right)
            start >>= 1
            stop >>= 1
        return self._func(res_left, res_right)

    def __repr__(self):
        return 'SegmentTree({0})'.format(self.data)
t = int(stdin.readline())
for _ in range(t):
    n = int(stdin.readline())
    a = [int(x) for x in stdin.readline().split()]
    prefix = [0]
    for i in range(n):
        prefix.append(prefix[-1] + a[i])
    prefix2 = [(prefix[i], -i) for i in range(n + 1)]
    prefix2.sort()
    prefix_address = {}
    for j in range(n + 1):
        prefix_address[-prefix2[j][1]] = j
    segtree = SegmentTree([-10 ** 6] * (n + 1))
    segtree[prefix_address[0]] = 1
    prev_dp = 0
    for i in range(n):
        if a[i] < 0:
            next_dp = -1
        elif a[i] == 0:
            next_dp = 0
        else:
            next_dp = 1
        next_dp += prev_dp
        next_dp = max(next_dp, segtree.query(0, prefix_address[i + 1]) + i)
        segtree[prefix_address[i + 1]] = next_dp - i
        prev_dp = next_dp
    stdout.write(str(next_dp) + '\n')