Lemper Cooking Competition (Codeforces 1725L)
Tags: Others, data structures

Pak Chanek is participating in a lemper cooking competition. In the competition, Pak Chanek has to cook lempers with $$$N$$$ stoves that are arranged sequentially from stove $$$1$$$ to stove $$$N$$$. Initially, stove $$$i$$$ has a temperature of $$$A_i$$$ degrees. A stove can have a negative temperature.Pak Chanek realises that, in order for his lempers to be cooked, he needs to keep the temperature of each stove at a non-negative value. To make it happen, Pak Chanek can do zero or more operations. In one operation, Pak Chanek chooses one stove $$$i$$$ with $$$2 \leq i \leq N-1$$$, then:  changes the temperature of stove $$$i-1$$$ into $$$A_{i-1} := A_{i-1} + A_{i}$$$,  changes the temperature of stove $$$i+1$$$ into $$$A_{i+1} := A_{i+1} + A_{i}$$$, and  changes the temperature of stove $$$i$$$ into $$$A_i := -A_i$$$. Pak Chanek wants to know the minimum number of operations he needs to do such that the temperatures of all stoves are at non-negative values. Help Pak Chanek by telling him the minimum number of operations needed or by reporting if it is not possible to do.

Input: The first line contains a single integer $$$N$$$ ($$$1 \le N \le 10^5$$$) — the number of stoves.The second line contains $$$N$$$ integers $$$A_1, A_2, \ldots, A_N$$$ ($$$-10^9 \leq A_i \leq 10^9$$$) — the initial temperatures of the stoves.

Output: Output an integer representing the minimum number of operations needed to make the temperatures of all stoves at non-negative values or output $$$-1$$$ if it is not possible.

Note: For the first example, a sequence of operations that can be done is as follows:   Pak Chanek does an operation to stove $$$3$$$, $$$A = [2, -2, 1, 4, 2, -2, 9]$$$.  Pak Chanek does an operation to stove $$$2$$$, $$$A = [0, 2, -1, 4, 2, -2, 9]$$$.  Pak Chanek does an operation to stove $$$3$$$, $$$A = [0, 1, 1, 3, 2, -2, 9]$$$.  Pak Chanek does an operation to stove $$$6$$$, $$$A = [0, 1, 1, 3, 0, 2, 7]$$$. There is no other sequence of operations such that the number of operations needed is fewer than $$$4$$$.

Sample Inputs:
7
2 -1 -1 5 2 -2 9

5
-1 -2 -3 -4 -5

Sample Outputs:
4

-1

Solution (Python):
import sys
input = sys.stdin.buffer.readline

class SegmentTree:

    def __init__(self, data, default=0, func=lambda x, y: x + y):
        """"""
        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)

def process(A):
    n = len(A)
    A1 = []
    curr = 0
    for i in range(n):
        curr += A[i]
        A1.append([curr, i])
    if min(A1)[0] < 0 or max(A1)[0] != A1[-1][0]:
        print('-1')
        return
    A1.sort()
    A1 = [[i, A1[i][1]] for i in range(n)]
    A1.sort(key=lambda a: a[1])
    S = SegmentTree(data=[0 for i in range(n)])
    answer = 0
    for (x, i) in A1:
        answer += S.query(x + 1, n)
        S[x] = 1
    print(answer)
n = int(input())
A = [int(x) for x in input().split()]
process(A)

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

def debug(*args):
    if os.environ.get('debug') in [None, '0']:
        return
    from inspect import currentframe, getframeinfo
    from re import search
    frame = currentframe().f_back
    s = getframeinfo(frame).code_context[0]
    r = search('\\((.*)\\)', s).group(1)
    vnames = r.split(', ')
    var_and_vals = [f'{var}={val}' for (var, val) in zip(vnames, args)]
    prefix = f'{currentframe().f_back.f_lineno:02d}: '
    print(f"{prefix}{', '.join(var_and_vals)}")
from types import GeneratorType

def bootstrap(f, stack=[]):

    def wrappedfunc(*args, **kwargs):
        if stack:
            return f(*args, **kwargs)
        to = f(*args, **kwargs)
        while True:
            if type(to) is GeneratorType:
                stack.append(to)
                to = next(to)
            else:
                stack.pop()
                if not stack:
                    break
                to = stack[-1].send(to)
        return to
    return wrappedfunc

class IntKeyDict(dict):
    from random import randrange
    rand = randrange(1 << 62)

    def __setitem__(self, k, v):
        super().__setitem__(k ^ self.rand, v)

    def __getitem__(self, k):
        return super().__getitem__(k ^ self.rand)

    def __contains__(self, k):
        return super().__contains__(k ^ self.rand)

    def __repr__(self):
        return str({k: v for (k, v) in self.items()})

    def get(self, k, default=None):
        return super().get(k ^ self.rand, default)

    def keys(self):
        return [k ^ self.rand for k in super().keys()]

    def items(self):
        return [(k ^ self.rand, v) for (k, v) in super().items()]
INF = float('inf')

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 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 _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 __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))

    def val(self, it):
        """"""
        (pos, idx) = it
        return self._lists[pos][idx]

    def begin(self):
        """"""
        return (0, 0)

    def end(self):
        """"""
        return (len(self._lists) - 1, len(self._lists[-1])) if self._lists else (0, 0)

    def prev(self, it):
        """"""
        if it == self.begin():
            raise ValueError('{0!r} already list begin'.format(it))
        (pos, idx) = it
        return (pos, idx - 1) if idx else (pos - 1, len(self._lists[pos - 1]) - 1)

    def next(self, it):
        """"""
        if it == self.end():
            raise ValueError('{0!r} already list end'.format(it))
        (pos, idx) = it
        return (pos, idx + 1) if pos + 1 == len(self._lists) or idx + 1 != len(self._lists[pos]) else (pos + 1, 0)
from itertools import accumulate

def main():
    N = int(input())
    A = list(map(int, input().split()))
    A = list(accumulate(A))
    if min(A) < 0 or A[-1] != max(A):
        return -1
    res = 0
    sl = SortedList()
    for a in A:
        res += len(sl) - sl.bisect_right(a)
        sl.add(a)
    return res
if __name__ == '__main__':
    res = main()
    print(res)

Solution (Python):
(standard_input, packages, output_together) = (1, 1, 1)
(dfs, hashing, read_from_file) = (0, 0, 0)
de = 1
if 1:
    if standard_input:
        import io, os, sys
        input = lambda : sys.stdin.readline().strip()
        import math
        inf = math.inf

        def I():
            return input()

        def II():
            return int(input())

        def MII():
            return map(int, input().split())

        def LI():
            return list(input().split())

        def LII():
            return list(map(int, input().split()))

        def LFI():
            return list(map(float, input().split()))

        def GMI():
            return map(lambda x: int(x) - 1, input().split())

        def LGMI():
            return list(map(lambda x: int(x) - 1, input().split()))
    if packages:
        from io import BytesIO, IOBase
        import random
        import os
        import bisect
        import typing
        from collections import Counter, defaultdict, deque
        from copy import deepcopy
        from functools import cmp_to_key, lru_cache, reduce
        from heapq import merge, heapify, heappop, heappush, heappushpop, nlargest, nsmallest
        from itertools import accumulate, combinations, permutations, count, product
        from operator import add, iand, ior, itemgetter, mul, xor
        from string import ascii_lowercase, ascii_uppercase, ascii_letters
        from typing import *
        BUFSIZE = 4096
    if output_together:

        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.stdout = IOWrapper(sys.stdout)
    if dfs:
        from types import GeneratorType

        def bootstrap(f, stack=[]):

            def wrappedfunc(*args, **kwargs):
                if stack:
                    return f(*args, **kwargs)
                else:
                    to = f(*args, **kwargs)
                    while True:
                        if type(to) is GeneratorType:
                            stack.append(to)
                            to = next(to)
                        else:
                            stack.pop()
                            if not stack:
                                break
                            to = stack[-1].send(to)
                    return to
            return wrappedfunc
    if hashing:
        RANDOM = random.getrandbits(20)

        class Wrapper(int):

            def __init__(self, x):
                int.__init__(x)

            def __hash__(self):
                return super(Wrapper, self).__hash__() ^ RANDOM
    if read_from_file:
        file = open('input.txt', 'r').readline().strip()[1:-1]
        fin = open(file, 'r')
        input = lambda : fin.readline().strip()
        output_file = open('output.txt', 'w')

        def fprint(*args, **kwargs):
            print(*args, **kwargs, file=output_file)
    if de:

        def debug(*args, **kwargs):
            print('\x1b[92m', end='')
            print(*args, **kwargs)
            print('\x1b[0m', end='')

class Solution:

    def reversePairs(self, nums: List[int]) -> int:
        ans = 0

        def getReverse(lst):
            nonlocal ans
            if len(lst) < 2:
                return lst
            half = len(lst) // 2
            half_ = len(lst) - len(lst) // 2
            (first_half, second_half) = (getReverse(lst[:half]), getReverse(lst[half:]))
            lst = []
            i1 = 0
            i2 = 0
            while i1 < half or i2 < half_:
                if i1 == half:
                    lst += second_half[i2:]
                    break
                if i2 == half_:
                    lst += first_half[i1:]
                    ans += (half - i1) * i2
                    break
                if first_half[i1] <= second_half[i2]:
                    lst.append(first_half[i1])
                    i1 += 1
                    ans += i2
                else:
                    lst.append(second_half[i2])
                    i2 += 1
            return lst
        getReverse(nums)
        return ans

def main():
    n = II()
    nums = LII()
    for i in range(1, n):
        nums[i] += nums[i - 1]
    if min(nums) < 0 or max(nums) != nums[-1]:
        print(-1)
        return
    print(Solution().reversePairs(nums))
    return
t = 1
for _ in range(t):
    main()