Points (Codeforces 1701F)
Tags: Math, matrices

A triple of points $$$i$$$, $$$j$$$ and $$$k$$$ on a coordinate line is called beautiful if $$$i < j < k$$$ and $$$k - i \le d$$$.You are given a set of points on a coordinate line, initially empty. You have to process queries of three types:  add a point;  remove a point;  calculate the number of beautiful triples consisting of points belonging to the set. 

Input: The first line contains two integers $$$q$$$ and $$$d$$$ ($$$1 \le q, d \le 2 \cdot 10^5$$$) — the number of queries and the parameter for defining if a triple is beautiful, respectively.The second line contains $$$q$$$ integers $$$a_1, a_2, \dots, a_q$$$ ($$$1 \le a_i \le 2 \cdot 10^5$$$) denoting the queries. The integer $$$a_i$$$ denotes the $$$i$$$-th query in the following way:  if the point $$$a_i$$$ belongs to the set, remove it; otherwise, add it;  after adding or removing the point, print the number of beautiful triples. 

Output: For each query, print one integer — the number of beautiful triples after processing the respective query.


Sample Inputs:
7 5
8 5 3 2 1 5 6

Sample Outputs:
0
0
1
2
5
1
5

Solution (Python):
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')
MAXN = 200000

class SegmentTree:

    def __init__(self):
        self.lazy = []
        self.v0 = []
        self.v1 = []
        self.v2 = []
        self.active = []
        for _ in range(4 * MAXN + 1):
            self.lazy.append(0)
            self.v0.append(1)
            self.v1.append(0)
            self.v2.append(0)
            self.active.append(True)
        self._init_active(1, 1, MAXN)

    def _init_active(self, x, l, r):
        if l == r:
            self.active[x] = False
        else:
            m = (l + r) // 2
            self._init_active(x * 2, l, m)
            self._init_active(x * 2 + 1, m + 1, r)
            self._reclac(x)

    def _reclac(self, x):
        self.v0[x] = 0
        self.v1[x] = 0
        self.v2[x] = 0
        if self.active[x * 2]:
            self.v0[x] += self.v0[x * 2]
            self.v1[x] += self.v1[x * 2]
            self.v2[x] += self.v2[x * 2]
        if self.active[x * 2 + 1]:
            self.v0[x] += self.v0[x * 2 + 1]
            self.v1[x] += self.v1[x * 2 + 1]
            self.v2[x] += self.v2[x * 2 + 1]

    def _push(self, x, l, r):
        la = self.lazy[x]
        if la == 0:
            return
        if l != r:
            self.v2[x * 2] = self.v2[x * 2] + 2 * la * self.v1[x * 2] + la * la * self.v0[x * 2]
            self.v1[x * 2] += la * self.v0[x * 2]
            self.v2[x * 2 + 1] = self.v2[x * 2 + 1] + 2 * la * self.v1[x * 2 + 1] + la * la * self.v0[x * 2 + 1]
            self.v1[x * 2 + 1] += la * self.v0[x * 2 + 1]
            self.lazy[x * 2] += la
            self.lazy[x * 2 + 1] += la
        self.lazy[x] = 0

    def update(self, x, l, r, a, b, up: bool):
        if r < a or b < l:
            return
        if a <= l and r <= b:
            if up:
                self.v2[x] = self.v2[x] + 2 * self.v1[x] + self.v0[x]
                self.v1[x] += self.v0[x]
                self.lazy[x] += 1
            else:
                self.v2[x] = self.v2[x] - 2 * self.v1[x] + self.v0[x]
                self.v1[x] -= self.v0[x]
                self.lazy[x] -= 1
            return
        m = (l + r) // 2
        self._push(x, l, r)
        self.update(x * 2, l, m, a, b, up)
        self.update(x * 2 + 1, m + 1, r, a, b, up)
        self._reclac(x)

    def set_state(self, x, l, r, pos, up: bool):
        if pos < l or r < pos:
            return
        if l == r:
            self.active[x] = up
            return
        m = (l + r) // 2
        self._push(x, l, r)
        self.set_state(x * 2, l, m, pos, up)
        self.set_state(x * 2 + 1, m + 1, r, pos, up)
        self._reclac(x)

def solve():
    (q, d) = map(int, input().split())
    points = map(int, input().split())
    tree = SegmentTree()
    check = [0] * (MAXN + 1)
    ans = []
    for point in points:
        if check[point]:
            check[point] = 0
            tree.update(1, 1, MAXN, max(1, point - d), point - 1, False)
            tree.set_state(1, 1, MAXN, point, False)
        else:
            check[point] = 1
            tree.update(1, 1, MAXN, max(1, point - d), point - 1, True)
            tree.set_state(1, 1, MAXN, point, True)
        v1 = tree.v1[1]
        v2 = tree.v2[1]
        ans.append((v2 - v1) // 2)
    print('\n'.join(map(str, ans)))
if __name__ == '__main__':
    solve()

Solution (Python):
(standard_input, packages, output_together) = (1, 1, 0)
(dfs, hashing, read_from_file) = (1, 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 LazySegTree:

    def __init__(self, n: int) -> None:
        self._n = n
        self._log = (self._n - 1).bit_length()
        self._size = 1 << self._log
        self._alive = [0] * (2 * self._size)
        self._margin = [0] * (2 * self._size)
        self._ans = [0] * (2 * self._size)
        self._lz = [0] * self._size

    def set(self, p: int, x: typing.Any) -> None:
        assert 0 <= p < self._n
        p += self._size
        for i in range(self._log, 0, -1):
            self._push(p >> i)
        if x == -1:
            self._alive[p] = 0
            self._margin[p] = 0
            self._ans[p] = 0
        else:
            self._alive[p] = 1
            self._margin[p] = x
            self._ans[p] = x * (x - 1) // 2
        for i in range(1, self._log + 1):
            self._update(p >> i)

    def all_prod(self) -> typing.Any:
        return self._ans[1]

    def apply(self, left: int, right: typing.Optional[int]=None, f: typing.Optional[typing.Any]=None) -> None:
        assert f is not None
        if right is None:
            p = left
            assert 0 <= left < self._n
            p += self._size
            for i in range(self._log, 0, -1):
                self._push(p >> i)
            self._d[p] = self._mapping(f, self._d[p])
            for i in range(1, self._log + 1):
                self._update(p >> i)
        else:
            assert 0 <= left <= right <= self._n
            if left == right:
                return
            left += self._size
            right += self._size
            for i in range(self._log, 0, -1):
                if left >> i << i != left:
                    self._push(left >> i)
                if right >> i << i != right:
                    self._push(right - 1 >> i)
            l2 = left
            r2 = right
            while left < right:
                if left & 1:
                    self._all_apply(left, f)
                    left += 1
                if right & 1:
                    right -= 1
                    self._all_apply(right, f)
                left >>= 1
                right >>= 1
            left = l2
            right = r2
            for i in range(1, self._log + 1):
                if left >> i << i != left:
                    self._update(left >> i)
                if right >> i << i != right:
                    self._update(right - 1 >> i)

    def _update(self, k: int) -> None:
        self._alive[k] = self._alive[k * 2] + self._alive[k * 2 + 1]
        self._margin[k] = self._margin[k * 2] + self._margin[k * 2 + 1]
        self._ans[k] = self._ans[k * 2] + self._ans[k * 2 + 1]

    def _all_apply(self, k: int, f: typing.Any) -> None:
        self._ans[k] += (self._margin[k] * 2 + (f - 1) * self._alive[k]) * f // 2
        self._margin[k] += self._alive[k] * f
        if k < self._size:
            self._lz[k] = f + self._lz[k]

    def _push(self, k: int) -> None:
        self._all_apply(2 * k, self._lz[k])
        self._all_apply(2 * k + 1, self._lz[k])
        self._lz[k] = 0

class FenwickTree:

    def __init__(self, n):
        self.n = n
        self.bit = [0] * n

    def sum(self, r):
        res = 0
        while r >= 0:
            res += self.bit[r]
            r = (r & r + 1) - 1
        return res

    def rsum(self, l, r):
        return self.sum(r) - self.sum(l - 1)

    def add(self, idx, delta):
        while idx < self.n:
            self.bit[idx] += delta
            idx = idx | idx + 1
(q, d) = MII()
queries = LGMI()
M = 10 ** 5 * 2
fen = FenwickTree(M)
seg = LazySegTree(M)
status = [0] * M
ans = []
for v in queries:
    if status[v] == 1:
        seg.set(v, -1)
        seg.apply(v + 1, min(v + d + 1, M), -1)
        fen.add(v, -1)
    else:
        seg.set(v, fen.rsum(max(0, v - d), v))
        seg.apply(v + 1, min(v + d + 1, M), 1)
        fen.add(v, 1)
    status[v] ^= 1
    ans.append(seg.all_prod())
print('\n'.join(map(str, ans)))

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

def fenwick_tree(n):
    tree = [0] * (n + 1)
    return tree

def get_sum0(i):
    s = 0
    while i > 0:
        s += tree[i]
        i -= i & -i
    return s

def get_sum(s, t):
    ans = 0 if s > t else get_sum0(t) - get_sum0(s - 1)
    return ans

def add(i, x):
    while i < len(tree):
        tree[i] += x
        i += i & -i

def f(i):
    if not lazy[i]:
        return
    li = lazy[i]
    for j in range(i * m, (i + 1) * m):
        if now[j]:
            x[j] += li
    lazy[i] = 0
    return
(p, d) = map(int, input().split())
a = list(map(int, input().split()))
m = 369
n = (2 * pow(10, 5) + m) // m + 1
tree = fenwick_tree(n * m)
now = [0] * (n * m)
cnt = [0] * n
(x, y) = ([0] * (n * m), [0] * n)
lazy = [0] * n
ans = []
ans0 = 0
for i in a:
    (l, r) = (i + 1, min(n * m - 1, i + d))
    c0 = get_sum(l, r)
    c = c0 * (c0 - 1) // 2
    (l, r) = (max(0, i - d), i - 1)
    (l0, r0) = (l // m, r // m)
    f(l0)
    f(i // m)
    if not now[i]:
        y[i // m] += c0
        cnt[i // m] += 1
        x[i] = c0
        while l <= r and l % m:
            if now[l]:
                c += x[l]
                x[l] += 1
                y[l0] += 1
            l += 1
        l0 = l // m
        while l + m - 1 <= r:
            c += y[l0]
            y[l0] += cnt[l0]
            lazy[l0] += 1
            l0 += 1
            l += m
        while l <= r:
            if now[l]:
                c += x[l]
                x[l] += 1
                y[l0] += 1
            l += 1
        ans0 += c
        add(i, 1)
    else:
        y[i // m] -= c0
        cnt[i // m] -= 1
        x[i] = 0
        while l <= r and l % m:
            if now[l]:
                x[l] -= 1
                y[l0] -= 1
                c += x[l]
            l += 1
        l0 = l // m
        while l + m - 1 <= r:
            y[l0] -= cnt[l0]
            lazy[l0] -= 1
            c += y[l0]
            l0 += 1
            l += m
        while l <= r:
            if now[l]:
                x[l] -= 1
                y[l0] -= 1
                c += x[l]
            l += 1
        ans0 -= c
        add(i, -1)
    now[i] ^= 1
    ans.append(ans0)
sys.stdout.write('\n'.join(map(str, ans)))