ABC170 E問題 Smart Infants

内容

• ABC170のE問題(500) Smart Infantsを解いた。
• コンテスト中には解けなかったが、この類の問題を解けるようになりたいので、メモで残す。

問題

• 園児と幼稚園がある。
• 園児は、それぞれどこかの幼稚園に所属していて、atcoderのrateを持っている。 *園児は、転園をする。(query)
• 各queryの後に、各幼稚園の(所属園児のrateの最大値)を求め、それの最小値を求める。

atcoder.jp

ポイント

• 各幼稚園はheapqを持っておいて、最大のrateを求めれるようにする。
• Seg木を作って、各幼稚園の最大rateをノードとしてもち、最大rateの最小値を求めれるようにする。
• 高速化のために、各園児の現在の所属する幼稚園を記録しておくと良い。

難しかったとこ

• 転園によって、園児が出ていく幼稚園と、入ってくる幼稚園がある。
• 入ってくる幼稚園は、heappushするだけなので、簡単
• 出ていく方は、heapqから要素を取り除く必要がある。
  • それが最大値ならheappopで良いが、そうでないとheapqの途中の値を取り除く必要が出てきて、それは辛い。

    ポイント

• 実際に取り除く処理を毎回するのではなく、サボって必要になった時にやる!!!

実装

import heapq
from collections import defaultdict

class SegmentTree():
    '''
    非再帰
    segment tree
    '''

    def __init__(self, n, func, init=float('inf')):
        '''
        n->配列の長さ
        func:func(a,b)->val,　func=minだとRMQになる
        木の高さhとすると,
        n:h-1までのノード数。h段目のノードにアクセスするために使う。
        data:ノード。
        parent:k->child k*2+1とk*2+2
        '''
        self.n = 2**(n-1).bit_length()
        self.init = init
        self.data = [init]*(2*self.n)
        self.func = func

    def set(self, k, v):
        '''
        あたいの初期化
        '''
        self.data[k+self.n-1] = v

    def build(self):
        '''
        setの後に一斉更新
        '''
        for k in reversed(range(self.n-1)):
            self.data[k] = self.func(self.data[k*2+1], self.data[k*2+2])

    def update(self, k, a):
        '''
        list[k]=aに更新する。
        rootまで更新
        '''
        k += self.n-1
        self.data[k] = a

        while k > 0:
            k = (k-1)//2
            self.data[k] = self.func(self.data[k*2+1], self.data[k*2+2])

    def query(self, l, r):
        '''
        [l,r)のfuncを求める
        '''
        L = l+self.n
        R = r+self.n
        ret = self.init
        while L < R:
            if R & 1:
                R -= 1
                ret = self.func(ret, self.data[R-1])
            if L & 1:
                ret = self.func(ret, self.data[L-1])
                L += 1
            L >>= 1
            R >>= 1
        return ret


############
#入力の処理
############
N,Q = map(int,input().split())

infants=[] #各園児の初期状態 kinder:(kindergarden)幼稚園の意
for _ in range(N):
    rate,kinder_id = map(int,input().split())
    infants.append((rate,kinder_id-1))

query = []
for _ in range(Q):
    child_id, rate = map(lambda x:int(x)-1,input().split())
    query.append((child_id,rate))

##############
#初期状態を作る
##############
MAXV=2*(10**5)+1 #幼稚園の最大数

belong = [-1]*N #各園児の所属 (kinder_id,rate)
kinder = defaultdict(list) #各園についての、rateのheapq(-rate,child_id)
maxs = [float("inf")]*MAXV #各園の最強園児のrate

#初期値
for child_id,(rate,kinder_id) in enumerate(infants):
    belong[child_id]=(kinder_id,rate)
    heapq.heappush(kinder[kinder_id],(-rate,child_id)) #maxが欲しいからマイナスをつけてheapqに入れる

#seg木を作る
Seg = SegmentTree(MAXV,min,float("inf"))
for kinder_id in kinder.keys():
    (rate,child_id) = kinder[kinder_id][0] #kinder[kinder_id][0]は、幼稚園kinder_idの最強園児
    Seg.set(kinder_id,-rate)
Seg.build()

##############
#queryの処理
#############
for child_id,to in query:
    #所属の更新
    (from_,rate) = belong[child_id]
    belong[child_id] = (to,rate)

    #転園の処理
    if kinder[from_][0][1]==child_id:#転園する奴が、その園の最強園児なら２番目に強い奴が新しく最強園児になる。
        heapq.heappop(kinder[from_])#転園する人を消す

        #新しい最強を見つける
        while kinder[from_]: 
            saikyo_child_id = kinder[from_][0][1]
            if belong[saikyo_child_id][0]==from_: #もしそいつが、from_に所属しているなら、そいつはfrom_で最強
                break
            else:#他の園のやつ(転園処理をさぼったから残ってるやつ)ならremove
                heapq.heappop(kinder[from_])

        if kinder[from_]: #kinder[frmo_]の先頭には、新しく最強園児になる奴がきてる。
            Seg.update(from_, -kinder[from_][0][0])
        else:#誰もいない場合
            Seg.update(from_, float("inf"))
    else:
        pass#転園するのが最強園児でないなら、segを更新しなくていいから、何もしない。heapqからのremoveはサボる。(あとでやる)

    #入園の処理
    heapq.heappush(kinder[to],(-rate,child_id))
    Seg.update(to,-kinder[to][0][0])

    #segに各園の最大値が入っているから、その最小値を求める
    print(Seg.query(0,MAXV+1))

終わりに

• heapqから、先頭以外の要素を取り除きたい気持ちになったときの解決法を学べた気がする。
  • 普通に削除するのはpythonだとできない気がしてて、後から処理する方法があるとしれてよかった。