Game with Stones (Codeforces 1589E)
Tags: Greedy, games

Bob decided to take a break from calculus homework and designed a game for himself. The game is played on a sequence of piles of stones, which can be described with a sequence of integers $$$s_1, \ldots, s_k$$$, where $$$s_i$$$ is the number of stones in the $$$i$$$-th pile. On each turn, Bob picks a pair of non-empty adjacent piles $$$i$$$ and $$$i+1$$$ and takes one stone from each. If a pile becomes empty, its adjacent piles do not become adjacent. The game ends when Bob can't make turns anymore. Bob considers himself a winner if at the end all piles are empty.We consider a sequence of piles winning if Bob can start with it and win with some sequence of moves.You are given a sequence $$$a_1, \ldots, a_n$$$, count the number of subsegments of $$$a$$$ that describe a winning sequence of piles. In other words find the number of segments $$$[l, r]$$$ ($$$1 \leq l \leq r \leq n$$$), such that the sequence $$$a_l, a_{l+1}, \ldots, a_r$$$ is winning.

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

Output: Print a single integer for each test case — the answer to the problem.

Note: In the first test case, Bob can't win on subsegments of length $$$1$$$, as there is no pair of adjacent piles in an array of length $$$1$$$.In the second test case, every subsegment is not winning.In the fourth test case, the subsegment $$$[1, 4]$$$ is winning, because Bob can make moves with pairs of adjacent piles: $$$(2, 3)$$$, $$$(1, 2)$$$, $$$(3, 4)$$$. Another winning subsegment is $$$[2, 3]$$$.

Sample Inputs:
6
2
2 2
3
1 2 3
4
1 1 1 1
4
1 2 2 1
4
1 2 1 2
8
1 2 1 2 1 2 1 2

Sample Outputs:
1
0
4
2
1
3

Solution (Python):
import sys
input = sys.stdin.readline
from collections import deque

def process(a, n):
    b = [0] * n
    b[0] = a[0]
    for i in range(1, n):
        b[i] = a[i] - b[i - 1]
    return b

def check(a, b, n):
    c = deque([])
    c.append([0, 1])
    d = deque([])
    count = 0
    for i in range(n):
        if i % 2 == 0:
            while c and c[-1][0] > b[i]:
                c.pop()
            if c and c[-1][0] == b[i]:
                count += c[-1][1]
                c[-1][1] += 1
            else:
                c.append([b[i], 1])
            while d and b[i] < -d[0][0]:
                d.popleft()
            if d and b[i] == -d[0][0]:
                count += d[0][1]
        else:
            while c and b[i] < -c[0][0]:
                c.popleft()
            if c and b[i] == -c[0][0]:
                count += c[0][1]
            while d and d[-1][0] > b[i]:
                d.pop()
            if d and d[-1][0] == b[i]:
                count += d[-1][1]
                d[-1][1] += 1
            else:
                d.append([b[i], 1])
    return count
for _ in range(int(input())):
    n = int(input())
    a = [int(i) for i in input().split()]
    print(check(a, process(a, n), n))

Solution (Python):
import sys
input = sys.stdin.readline
from collections import deque

def process(a, n):
    b = [0] * n
    b[0] = a[0]
    for i in range(1, n):
        b[i] = a[i] - b[i - 1]
    return b

def check(a, b, n):
    c = deque([])
    c.append([0, 1])
    d = deque([])
    count = 0
    for i in range(n):
        if i % 2 == 0:
            while c and c[-1][0] > b[i]:
                c.pop()
            if c and c[-1][0] == b[i]:
                count += c[-1][1]
                c[-1][1] += 1
            else:
                c.append([b[i], 1])
            while d and b[i] < -d[0][0]:
                d.popleft()
            if d and b[i] == -d[0][0]:
                count += d[0][1]
        else:
            while c and b[i] < -c[0][0]:
                c.popleft()
            if c and b[i] == -c[0][0]:
                count += c[0][1]
            while d and d[-1][0] > b[i]:
                d.pop()
            if d and d[-1][0] == b[i]:
                count += d[-1][1]
                d[-1][1] += 1
            else:
                d.append([b[i], 1])
    return count
for _ in range(int(input())):
    n = int(input())
    a = [int(i) for i in input().split()]
    print(check(a, process(a, n), n))

Solution (Python):
from sys import stdin, stdout

class SegmentTree:

    def __init__(self, data, default=10 ** 18, func=min):
        """"""
        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 = [0] + [int(x) for x in stdin.readline().split()]
    if n == 1:
        if a[1] == 0:
            stdout.write('1\n')
        else:
            stdout.write('0\n')
    else:
        b = [0]
        dict1 = {}
        dict2 = {}
        for i in range(1, n + 1):
            b.append(a[i] - b[-1])
            if i % 2 == 0:
                if b[-1] in dict2:
                    dict2[b[-1]].append(i)
                else:
                    dict2[b[-1]] = [i]
            elif b[-1] in dict1:
                dict1[b[-1]].append(i)
            else:
                dict1[b[-1]] = [i]
        tree1 = SegmentTree(b[1::2])
        tree2 = SegmentTree(b[2::2])
        answer = 0
        for l in range(1, n + 1):
            if l % 2 == 1:
                if tree1.data[1] >= -b[l - 1]:
                    max_r_1 = n + 1
                else:
                    pointer = 1
                    while pointer < tree1._size:
                        if tree1.data[2 * pointer] < -b[l - 1]:
                            pointer *= 2
                        else:
                            pointer = 2 * pointer + 1
                    max_r_1 = 2 * (pointer - tree1._size) + 1
                if tree2.data[1] >= b[l - 1]:
                    max_r_2 = n + 1
                else:
                    pointer = 1
                    while pointer < tree2._size:
                        if tree2.data[2 * pointer] < b[l - 1]:
                            pointer *= 2
                        else:
                            pointer = 2 * pointer + 1
                    max_r_2 = 2 * (pointer - tree2._size + 1)
                max_r = min(max_r_1, max_r_2)
                tree1[l // 2] = 10 ** 18
                if -b[l - 1] in dict1:
                    to_search = dict1[-b[l - 1]]
                    if to_search[-1] >= l and to_search[0] < max_r:
                        lower = -1
                        upper = len(to_search) - 1
                        while upper - lower > 1:
                            candidate = (lower + upper) // 2
                            if to_search[candidate] >= l:
                                upper = candidate
                            else:
                                lower = candidate
                        l_pointer = upper
                        lower = 0
                        upper = len(to_search)
                        while upper - lower > 1:
                            candidate = (lower + upper) // 2
                            if to_search[candidate] >= max_r:
                                upper = candidate
                            else:
                                lower = candidate
                        r_pointer = lower
                        answer += r_pointer - l_pointer + 1
                if b[l - 1] in dict2:
                    to_search = dict2[b[l - 1]]
                    if to_search[-1] >= l and to_search[0] < max_r:
                        lower = -1
                        upper = len(to_search) - 1
                        while upper - lower > 1:
                            candidate = (lower + upper) // 2
                            if to_search[candidate] >= l:
                                upper = candidate
                            else:
                                lower = candidate
                        l_pointer = upper
                        lower = 0
                        upper = len(to_search)
                        while upper - lower > 1:
                            candidate = (lower + upper) // 2
                            if to_search[candidate] >= max_r:
                                upper = candidate
                            else:
                                lower = candidate
                        r_pointer = lower
                        answer += r_pointer - l_pointer + 1
            else:
                if tree1.data[1] >= b[l - 1]:
                    max_r_1 = n + 1
                else:
                    pointer = 1
                    while pointer < tree1._size:
                        if tree1.data[2 * pointer] < b[l - 1]:
                            pointer *= 2
                        else:
                            pointer = 2 * pointer + 1
                    max_r_1 = 2 * (pointer - tree1._size) + 1
                if tree2.data[1] >= -b[l - 1]:
                    max_r_2 = n + 1
                else:
                    pointer = 1
                    while pointer < tree2._size:
                        if tree2.data[2 * pointer] < -b[l - 1]:
                            pointer *= 2
                        else:
                            pointer = 2 * pointer + 1
                    max_r_2 = 2 * (pointer - tree2._size + 1)
                max_r = min(max_r_1, max_r_2)
                tree2[l // 2 - 1] = 10 ** 18
                if b[l - 1] in dict1:
                    to_search = dict1[b[l - 1]]
                    if to_search[-1] >= l and to_search[0] < max_r:
                        lower = -1
                        upper = len(to_search) - 1
                        while upper - lower > 1:
                            candidate = (lower + upper) // 2
                            if to_search[candidate] >= l:
                                upper = candidate
                            else:
                                lower = candidate
                        l_pointer = upper
                        lower = 0
                        upper = len(to_search)
                        while upper - lower > 1:
                            candidate = (lower + upper) // 2
                            if to_search[candidate] >= max_r:
                                upper = candidate
                            else:
                                lower = candidate
                        r_pointer = lower
                        answer += r_pointer - l_pointer + 1
                if -b[l - 1] in dict2:
                    to_search = dict2[-b[l - 1]]
                    if to_search[-1] >= l and to_search[0] < max_r:
                        lower = -1
                        upper = len(to_search) - 1
                        while upper - lower > 1:
                            candidate = (lower + upper) // 2
                            if to_search[candidate] >= l:
                                upper = candidate
                            else:
                                lower = candidate
                        l_pointer = upper
                        lower = 0
                        upper = len(to_search)
                        while upper - lower > 1:
                            candidate = (lower + upper) // 2
                            if to_search[candidate] >= max_r:
                                upper = candidate
                            else:
                                lower = candidate
                        r_pointer = lower
                        answer += r_pointer - l_pointer + 1
        stdout.write(str(answer) + '\n')