Edit

【pythonでABC170を解説】E - Smart Infants

問題概要

問題ページ

問題文

AtCoder に参加している幼児が \(N\) 人おり、\(1\) から \(N\) の番号が付けられています。また、幼稚園が \(2\times 10^5\) 校あり、\(1\) から \(2\times 10^5\) の番号が付けられています。
幼児 \(i\) のレートは \(A_i\) であり、はじめ幼稚園 \(B_i\) に所属しています。

これから \(Q\) 回にわたって、転園が行われます。
\(j\) 回目の転園では、幼児 \(C_j\) の所属を幼稚園 \(D_j\) に変更します。

ここで、「平等さ」を、AtCoderに参加している幼児が一人以上いるような幼稚園それぞれについて園内で最もレートの高い幼児のレートを求め、その最小値として得られる値とします。

\(Q\) 回それぞれの転園後の平等さを求めてください。

制約

  • \(1 \leq N,Q \leq 2 \times 10^5\)
  • \(1 \leq A_i \leq 10^9\)
  • \(1 \leq C_j \leq N\)
  • \(1 \leq B_i,D_j \leq 2 \times 10^5\)
  • 入力はすべて整数である。
  • \(j\) 回目の転園の前後で幼児 \(C_j\) の所属は異なる。

問題の考察

ACコード

import sys
import collections
import heapq


class AlgSegmentTree:
    def __init__(self, size, f=lambda x, y: min(x, y), default=10 ** 15):
        self.size = 2**(size-1).bit_length()
        self.default = default
        self.dat = [default]*(self.size*2)
        self.f = f

    def update(self, i, x):
        i += self.size
        self.dat[i] = x
        while i > 0:
            i >>= 1
            self.dat[i] = self.f(self.dat[i*2], self.dat[i*2+1])

    def query(self, l, r):
        l += self.size
        r += self.size
        lres, rres = self.default, self.default
        while l < r:
            if l & 1:
                lres = self.f(lres, self.dat[l])
                l += 1

            if r & 1:
                r -= 1
                rres = self.f(self.dat[r], rres)
            l >>= 1
            r >>= 1
        res = self.f(lres, rres)
        return res


def if_not_exist(x, dic):
    if x not in dic:
        d = collections.defaultdict(int)
        hq = []
        heapq.heapify(hq)
        dic[x] += [d, hq]


def solve():
    input = sys.stdin.readline
    mod = 10 ** 9 + 7
    n, q = list(map(int, input().rstrip('\n').split()))
    st = AlgSegmentTree(n)
    dic = collections.defaultdict(list)
    ab = []
    st = AlgSegmentTree(2 * 10 ** 5)
    for i in range(n):
        a, b = list(map(int, input().rstrip('\n').split()))
        b -= 1
        if_not_exist(b, dic)
        ab.append([-a, b])
        dic[b][0][-a] += 1
        heapq.heappush(dic[b][1], -a)
    for k, v in dic.items():
        p = heapq.heappop(v[1])
        st.update(k, -p)
        heapq.heappush(v[1], p)
    for i in range(q):
        c, d = list(map(int, input().rstrip('\n').split()))
        c, d = c - 1, d - 1
        a, b = ab[c]
        if_not_exist(b, dic)
        dic[b][0][a] -= 1
        if_not_exist(d, dic)
        dic[d][0][a] += 1
        heapq.heappush(dic[d][1], a)
        #直前の所属の最小値を更新
        while True:
            if len(dic[b][1]):
                p = heapq.heappop(dic[b][1])
                if dic[b][0][p] != 0:
                    st.update(b, -p)
                    heapq.heappush(dic[b][1], p)
                    break
            else:
                st.update(b, 10 ** 15)
                break
        #現在の所属の最小値を更新
        while True:
            if len(dic[d][1]):
                p = heapq.heappop(dic[d][1])
                if dic[d][0][p] != 0:
                    st.update(d, -p)
                    heapq.heappush(dic[d][1], p)
                    break
            else:
                st.update(d, 0)
                break
        #直前の所属を更新
        ab[c][1] = d
        print(st.query(0, 2 * 10 ** 5))


if __name__ == '__main__':
    solve()

-Edit
-