Sum of Medians (Codeforces 85D)
Tags: Implement, binary search

In one well-known algorithm of finding the k-th order statistics we should divide all elements into groups of five consecutive elements and find the median of each five. A median is called the middle element of a sorted array (it's the third largest element for a group of five). To increase the algorithm's performance speed on a modern video card, you should be able to find a sum of medians in each five of the array.A sum of medians of a sorted k-element set S = {a1, a2, ..., ak}, where a1 < a2 < a3 < ... < ak, will be understood by as The  operator stands for taking the remainder, that is  stands for the remainder of dividing x by y.To organize exercise testing quickly calculating the sum of medians for a changing set was needed.

Input: The first line contains number n (1 ≤ n ≤ 105), the number of operations performed.Then each of n lines contains the description of one of the three operations:   add x — add the element x to the set;  del x — delete the element x from the set;  sum — find the sum of medians of the set. For any add x operation it is true that the element x is not included in the set directly before the operation.For any del x operation it is true that the element x is included in the set directly before the operation.All the numbers in the input are positive integers, not exceeding 109.

Output: For each operation sum print on the single line the sum of medians of the current set. If the set is empty, print 0.Please, do not use the %lld specificator to read or write 64-bit integers in C++. It is preferred to use the cin, cout streams (also you may use the %I64d specificator).


Sample Inputs:
6
add 4
add 5
add 1
add 2
add 3
sum

14
add 1
add 7
add 2
add 5
sum
add 6
add 8
add 9
add 3
add 4
add 10
sum
del 1
sum

Sample Outputs:
3

5
11
13

Solution (Python):
import io
import os
import heapq
from collections import Counter, deque
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')
ALPHA = 0.75

class ScapegoatTree:

    def __init__(self, key=0, deleted=True):
        self.left = None
        self.right = None
        self.siz = 0
        self.sm = [0 for _ in range(5)]
        self.key = key
        self.deleted = deleted
        if not self.deleted:
            self.siz += 1
            self.sm[0] += self.key

    def add(self, x):
        if x == self.key:
            if self.deleted:
                self.deleted = False
                self.__maintain()
                return True
            else:
                return False
        if x < self.key:
            if self.left is None:
                self.left = ScapegoatTree(key=x, deleted=False)
                self.__maintain()
                return True
            elif self.left.add(x):
                self.__maintain()
                if self.left.__unbalanced():
                    self.left = self.__rebuild(self.left)
                return True
        elif self.right is None:
            self.right = ScapegoatTree(key=x, deleted=False)
            self.__maintain()
            return True
        elif self.right.add(x):
            self.__maintain()
            if self.right.__unbalanced():
                self.right = self.__rebuild(self.right)
            return True
        return False

    def discard(self, x):
        if x == self.key:
            if self.deleted:
                return False
            else:
                self.deleted = True
                self.__maintain()
                return True
        flag = True
        if x < self.key:
            if self.left is None:
                flag = False
            else:
                flag = self.left.discard(x)
        elif self.right is None:
            flag = False
        else:
            flag = self.right.discard(x)
        if flag:
            self.__maintain()
        return flag

    def sum(self, i):
        return self.sm[i]

    def __maintain(self):
        for i in range(5):
            self.sm[i] = 0
        cnt = 0
        if self.left is not None:
            cnt += self.left.siz
            for i in range(5):
                self.sm[i] += self.left.sm[i]
        if not self.deleted:
            self.sm[cnt % 5] += self.key
            cnt += 1
        if self.right is not None:
            for i in range(5):
                self.sm[(cnt + i) % 5] += self.right.sm[i]
            cnt += self.right.siz
        self.siz = cnt

    def __unbalanced(self):
        if self is None or self.siz <= 5:
            return False
        if self.left is not None and self.left.siz > self.siz * ALPHA:
            return True
        if self.right is not None and self.right.siz > self.siz * ALPHA:
            return True
        return False

    @classmethod
    def __rebuild(cls, p):
        a = []

        def dfs(q):
            if q is None:
                return
            dfs(q.left)
            if not q.deleted:
                a.append(q)
            dfs(q.right)

        def build(l, r):
            if l > r:
                return None
            mid = (l + r) // 2
            a[mid].left = build(l, mid - 1)
            a[mid].right = build(mid + 1, r)
            a[mid].__maintain()
            return a[mid]
        dfs(p)
        return build(0, len(a) - 1)

def solve():
    n = int(input())
    tree = ScapegoatTree()
    for _ in range(n):
        line = input().strip()
        if line[:3] == 'sum':
            print(tree.sum(2))
        else:
            (opr, x) = [word for word in line.split()]
            x = int(x)
            if opr == 'add':
                tree.add(x)
            else:
                tree.discard(x)

def main():
    solve()
if __name__ == '__main__':
    main()

Solution (Python):
import io
import os
import heapq
from collections import Counter, deque
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')
ALPHA = 0.75
BETA = 7

class ScapegoatTree:

    def __init__(self, key=0, deleted=True):
        self.left = None
        self.right = None
        self.siz = 0
        self.sm = [0 for _ in range(5)]
        self.key = key
        self.deleted = deleted
        if not self.deleted:
            self.siz += 1
            self.sm[0] += self.key

    def add(self, x):
        if not self.__add(x):
            return False
        p = self
        while x != p.key:
            if x < p.key:
                if p.left.__unbalanced():
                    p.left = self.__rebuild(p.left)
                    break
                p = p.left
            else:
                if p.right.__unbalanced():
                    p.right = self.__rebuild(p.right)
                    break
                p = p.right
        return True

    def discard(self, x):
        if x == self.key:
            if self.deleted:
                return False
            else:
                self.deleted = True
                self.__maintain()
                return True
        flag = True
        if x < self.key:
            if self.left is None:
                flag = False
            else:
                flag = self.left.discard(x)
        elif self.right is None:
            flag = False
        else:
            flag = self.right.discard(x)
        if flag:
            self.__maintain()
        return flag

    def sum(self, i):
        return self.sm[i]

    def __add(self, x):
        if x == self.key:
            if self.deleted:
                self.deleted = False
                return True
            else:
                return False
        if x < self.key:
            if self.left is None:
                self.left = ScapegoatTree(key=x, deleted=False)
                flag = True
            else:
                flag = self.left.__add(x)
        elif self.right is None:
            self.right = ScapegoatTree(key=x, deleted=False)
            flag = True
        else:
            flag = self.right.__add(x)
        if flag:
            self.__maintain()
        return flag

    def __maintain(self):
        for i in range(5):
            self.sm[i] = 0
        cnt = 0
        if self.left is not None:
            cnt += self.left.siz
            for i in range(5):
                self.sm[i] += self.left.sm[i]
        if not self.deleted:
            self.sm[cnt % 5] += self.key
            cnt += 1
        if self.right is not None:
            for i in range(5):
                self.sm[(cnt + i) % 5] += self.right.sm[i]
            cnt += self.right.siz
        self.siz = cnt

    def __unbalanced(self):
        if self is None or self.siz <= BETA:
            return False
        if self.left is not None and self.left.siz > self.siz * ALPHA:
            return True
        if self.right is not None and self.right.siz > self.siz * ALPHA:
            return True
        return False

    @classmethod
    def __rebuild(cls, p):
        a = []

        def dfs(q):
            if q is None:
                return
            dfs(q.left)
            if not q.deleted:
                a.append(q)
            dfs(q.right)

        def build(l, r):
            if l > r:
                return None
            mid = (l + r) // 2
            a[mid].left = build(l, mid - 1)
            a[mid].right = build(mid + 1, r)
            a[mid].__maintain()
            return a[mid]
        dfs(p)
        return build(0, len(a) - 1)

def solve():
    n = int(input())
    tree = ScapegoatTree()
    for _ in range(n):
        line = input().strip()
        if line[:3] == 'sum':
            print(tree.sum(2))
        else:
            (opr, x) = [word for word in line.split()]
            x = int(x)
            if opr == 'add':
                tree.add(x)
            else:
                tree.discard(x)

def main():
    solve()
if __name__ == '__main__':
    main()

Solution (Python):
import io
import os
import heapq
from collections import Counter, deque
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')
ALPHA = 0.75
BETA = 15

class ScapegoatTree:

    def __init__(self, key=0, deleted=True):
        self.left = None
        self.right = None
        self.siz = 0
        self.sm = [0 for _ in range(5)]
        self.key = key
        self.deleted = deleted
        if not self.deleted:
            self.siz += 1
            self.sm[0] += self.key

    def add(self, x):
        if not self.__add(x):
            return False
        p = self
        while x != p.key:
            if x < p.key:
                if p.left.__unbalanced():
                    p.left = self.__rebuild(p.left)
                    break
                p = p.left
            else:
                if p.right.__unbalanced():
                    p.right = self.__rebuild(p.right)
                    break
                p = p.right
        return True

    def discard(self, x):
        if x == self.key:
            if self.deleted:
                return False
            else:
                self.deleted = True
                self.__maintain()
                return True
        flag = True
        if x < self.key:
            if self.left is None:
                flag = False
            else:
                flag = self.left.discard(x)
        elif self.right is None:
            flag = False
        else:
            flag = self.right.discard(x)
        if flag:
            self.__maintain()
        return flag

    def sum(self, i):
        return self.sm[i]

    def __add(self, x):
        if x == self.key:
            if self.deleted:
                self.deleted = False
                return True
            else:
                return False
        if x < self.key:
            if self.left is None:
                self.left = ScapegoatTree(key=x, deleted=False)
                flag = True
            else:
                flag = self.left.__add(x)
        elif self.right is None:
            self.right = ScapegoatTree(key=x, deleted=False)
            flag = True
        else:
            flag = self.right.__add(x)
        if flag:
            self.__maintain()
        return flag

    def __maintain(self):
        for i in range(5):
            self.sm[i] = 0
        cnt = 0
        if self.left is not None:
            cnt += self.left.siz
            for i in range(5):
                self.sm[i] += self.left.sm[i]
        if not self.deleted:
            self.sm[cnt % 5] += self.key
            cnt += 1
        if self.right is not None:
            for i in range(5):
                self.sm[(cnt + i) % 5] += self.right.sm[i]
            cnt += self.right.siz
        self.siz = cnt

    def __unbalanced(self):
        if self is None or self.siz <= BETA:
            return False
        if self.left is not None and self.left.siz > self.siz * ALPHA:
            return True
        if self.right is not None and self.right.siz > self.siz * ALPHA:
            return True
        return False

    @classmethod
    def __rebuild(cls, p):
        a = []

        def dfs(q):
            if q is None:
                return
            dfs(q.left)
            if not q.deleted:
                a.append(q)
            dfs(q.right)

        def build(l, r):
            if l > r:
                return None
            mid = (l + r) // 2
            a[mid].left = build(l, mid - 1)
            a[mid].right = build(mid + 1, r)
            a[mid].__maintain()
            return a[mid]
        dfs(p)
        return build(0, len(a) - 1)

def solve():
    n = int(input())
    tree = ScapegoatTree()
    for _ in range(n):
        line = input().strip()
        if line[:3] == 'sum':
            print(tree.sum(2))
        else:
            (opr, x) = [word for word in line.split()]
            x = int(x)
            if opr == 'add':
                tree.add(x)
            else:
                tree.discard(x)

def main():
    solve()
if __name__ == '__main__':
    main()