aacord’s memo

abcを中心にpythonで解いた問題のメモ、整理をしています。緑になった。

abc 127 F - Absolute Minima (python)

BIT木。初見で解いているときは BIT木を知らなくて TLE を出し続け、諦めて解説読んでナニコレとなった。
BIT木の使いどころは、BIT木を2本持つことで、小さい方から i 番目の値の探索を log(N) を行える点である。値の追加、探索、削除を log(N) でできるので、C++とかの set の代わりが出来る。
heap のやつとの違いは、BITは最小値以外の値も探索して出力できる点。heap は multiset にも対応できる点。だと思う。使い分けは実際に問題解いて理解するのが良さそう。
Fは求める x は | x - Ai | を最小にする x なので Ai の中央値となる。ついにこの知識を使う時が来たか。そして求める答えは

中央値以下の個数 * x - 中央値以下の総和 + 中央値以上の個数 * x - 中央値以上の総和 + C の総和となる。

なので、BIT 木で中央値以下の個数、中央値以下の総和、中央値以上の個数、中央値以上の総和を求める。
これを実現するために、値を管理するBITと要素数を管理するBITの二本のBIT木を持つ。
あと、値の取りうる範囲が 10**9 まであるので、クエリを先読みして、値の小さい順に番号をつけておいて、その番号でBIT木を管理する必要がある。要するに座圧。

import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline

class Bit:
    def __init__(self, n):
        self.size = n
        self.tree = [0] * (n + 1)
        self.depth = n.bit_length()
 
    def sum(self, i):
        s = 0
        while i > 0:
            s += self.tree[i]
            i -= i & -i
        return s
 
    def add(self, i, x):
        while i <= self.size:
            self.tree[i] += x
            i += i & -i
 
    def lower_bound(self, x):
        sum_ = 0
        pos = 0
        for i in range(self.depth, -1, -1):
            k = pos + (1 << i)
            if k <= self.size and sum_ + self.tree[k] < x:
                sum_ += self.tree[k]
                pos += 1 << i
        return pos + 1
  
def compress(list1):
    list2 = sorted(set(list1))
    memo = {value : index for index, value in enumerate(list2)}
    return memo, len(list2)

q = int(readline())
L = [list(map(int,readline().split())) for i in range(q)]
 
#座圧する
chk = []
for i in range(q):
    if L[i][0] == 1:
        chk.append(L[i][1])
        
memo, len_memo = compress(chk)
memo_inv = dict([(v,k) for k,v in memo.items()])

#値を管理するBIT
bit = Bit(len_memo)
#要素数を管理するBIT
num = Bit(len_memo)
  
s = 0
cnt = 0
for i in range(q):
  que = L[i]
  if que[0] == 1:
    a,b,c = que
    s += c
    cnt += 1
    bit.add(memo[b]+1,b)
    num.add(memo[b]+1,1)
  else:
    d = num.lower_bound((cnt+1)//2)
    m = memo_inv[d-1]
    sum1 = bit.sum(d)
    num1 = num.sum(d)
    sum2 = bit.sum(len_memo)-bit.sum(d)
    num2 = num.sum(len_memo)-num.sum(d)
    print(m,abs(num1*m-sum1)+abs(num2*m-sum2)+s)