AtCoder茶色解く5

今回解くのはこちら
D - Replacing

問題が複雑なので、実際の例をみて考えます。

まず、以下が入力で与えられます。

4
1 2 3 4
3
1 2
3 4
2 4

この時、配列1 2 3 4に対して、3回の操作を行います。
1を2に、3を4に、2を4にすると、それぞれ

2 2 3 4 = 11
2 2 4 4 = 12
4 4 4 4 = 16 

となります。

また、
配列の数、配列の取りうる数、取り替える数、変更値のすべてが10^5の範囲となります。
よって、2重ループでは時間が足りないので、何か工夫を考えます。

ここで、元の配列にそれぞれの数が何こずつ存在するかをハッシュで記憶しておくと、
そのハッシュ値を元に変更しながら計算できるのではと考えました。

流れとしては以下です。
・配列を取得
・配列の各値がいくつあるか値:個数でハッシュに格納、さらに合計値を計算しておく
・for文で、置換する値を受け取り、ハッシュの値を書き換える 書き換えた結果を元に再計算

これを繰り返していけばいいはず。
実装は以下です。

N=int(input())
A=list(map(int, input().split()))
l={}

# Aの数をまとめておく
for i in A:
  if i not in l:
    l[i]=1
  else:
    l[i]+=1
sum=sum(A)
#print(l, sum)

M=int(input())
for i in range(M):
  B, C = map(int, input().split())
  if B in l:
    bef=l[B]*B
    aft=l[B]*C
    if C in l:
      l[C]+=l[B]
    else:
      l[C]=l[B] 
    l[B]=0
    #print(l,sum, bef, aft)
    sum=sum-bef+aft
  print(sum)

サンプルでちょっとNGがあったので、処理を増やしています。
前回覚えた差分更新が早速役に立ちました。