AtCoderで公開されている、E869120さんの『競プロ典型90問』を解いていきます。
問題
長さ$N$の数列$A,\ B$について、$A$の各要素を1ずつ増やす/減らす操作を$K$回することで$A$を$B$に一致させることができるかを求める。
アイデア
$A,\ B$の各要素$a_i,\ b_i$について、その差の絶対値$|a_i - b_i|$が$a_i$を$b_i$に一致させるために必要な最小の操作数。
⇒$ K = \sum |a_i - b_i| $ならOK。
diff = 0
for a, b in A, B:
diff += abs(a - b)
でよさそう。
同じ要素$a_i$に対して、1増やして1減らせばもとに戻る。
⇒$K$は偶数回余分にあっても構わない！
改良案を考える
各要素の絶対値の和
for a in A:のノリでfor a, b in A, B:としてみたけど、だめだった。
⇒今更だけど、for ? in ???:の構文を真面目に勉強しよう。
for文の仕組み
-
for ? in A:は、ループの度にAを参照して次の要素を取り出す
A = [1, 2]
for a in A:
A.append(a)
print(a, A)
if len(A) > 5:
break
# 出力
# 1 [1, 2, 1]
# 2 [1, 2, 1, 2]
# 1 [1, 2, 1, 2, 1]
# 2 [1, 2, 1, 2, 1, 2]
つまり、ループ開始時にAを展開しておき、それを1つずつ参照していくわけではないということです。
# その場合予想される出力
# 1 [1, 2, 1]
# 2 [1, 2, 1, 2]
ただし、ループ中にAを書き換えるとRuntimeErrorが出ることがあり、基本的には避けるべきです。
-
for ? in A, B:のような場合、A, Bはtupleの(A, B)として解釈される
A = [1, 2]
B = [3, 4]
for a in A, B:
print(a)
# 出力
# [1, 2]
# [3, 4]
ただ単にa = A、a = Bとしてるだけです。もし、AとBの同じ番号の要素を取り出したいなら、zipを使います。
A = [1, 2]
B = [3, 4]
for a in zip(A, B):
print(a)
# 出力
# (1, 3)
# (2, 4)
なお、zipはtuple型のまとまりを返すため、list型とは異なり書き換え無効のようです。
A = [1, 2]
B = [3, 4]
for a in zip(A, B):
a[0] += 1
print(a)
# 期待する出力
# (2, 3)
# (3, 4)
# 実際の出力
# TypeError: 'tuple' object does not support item assignment
-
for a, b in ???:のような場合、???の各要素がイテレータで、かつ各要素の長さと変数の数（a, bの2つ）が一致している必要があります。a, bがtuple型の(a, b)として認識されることが理由のようです。
A = [(1, 2), (3, 4)]
for a, b in A:
print(a, b)
# 出力
# 1, 2
# 3, 4
問題に戻る
ということで、for a, b in A, B:でエラーが出る理由が分かりました（
- まずA, Bは(A, B)と認識され、for文はA, Bの順に参照
- A, Bはリストでいずれも長さ$N$であるため、$N \neq 2$のときに(a, b)の形で参照できない
）。
この問題では、for a, b in zip(A, B):を使うのがよさそうです。
最終的な実装例
N, K = map(int, input().split())
A = list(map(int, input().split()))
B = list(map(int, input().split()))
diff = 0
for a, b in zip(A, B):
diff += abs(a - b)
if diff - K <= 0 and (diff - K) % 2 == 0:
print("Yes")
else:
print("No")
一度for文についてじっくり考えてみました。まずは何となくコードを書けるようにして、謎のエラーに出くわしたら、随時深く考えていこうと思います。
お読みいただきありがとうございました。
参考記事