ABC171Dを解いた【個数管理+差分更新】

筆者はレート800前後の茶～緑コーダ

ABC171のD問題を解いていく

実装コード

同じ値が何個あるかをまとめて管理する。

そのために

• Counter を用いて、各値の出現回数を数え上げる
• 値の変更・加算・削除が発生するため、管理用の辞書には defaultdict を使う

各クエリでは

• 値 b の個数をまとめて取り出し
• それらを値 c に付け替える

という操作を行う。
総和は毎回すべてを計算し直す必要はなく、
個数 × 差分(c-b) だけを加減算すれば即座に更新できる。

main.py

from bisect import bisect_left, bisect_right, insort_left, insort_right
from collections import defaultdict, Counter, deque
from functools import reduce, lru_cache
from itertools import product, accumulate, groupby, combinations
import sys
import os
def rI(): return int(sys.stdin.readline().rstrip())
def rLI(): return list(map(int,sys.stdin.readline().rstrip().split()))
def rI1(): return (int(sys.stdin.readline().rstrip())-1)
def rLI1(): return list(map(lambda a:int(a)-1,sys.stdin.readline().rstrip().split()))
def rS(): return sys.stdin.readline().rstrip()
def rLS(): return list(sys.stdin.readline().rstrip().split())
IS_LOCAL = int(os.getenv("ATCODER", "0"))==0
err = (lambda *args, **kwargs: print(*args, **kwargs, file=sys.stderr)) if IS_LOCAL else (lambda *args, **kwargs: None)

def main():
    N = rI()
    A = rLI()
    Q = rI()
    d = defaultdict(lambda :0, Counter(A))
    S = sum(A)
    for _ in range(Q):
        b, c = rLI()
        x = d.pop(b, 0)
        d[c] += x
        S += x*(c-b)
        print(S)
        
if __name__ == '__main__':
    main()

感想

かなり素直な発想で、そのまま実装まで持っていけた問題だった

個人的お気に入りライブラリのdefaultdict,Counterの合せ技を使うことができてなんか嬉しい