ABC128-Bで詰まったので備忘録的に示しとく

公式ドキュメントが読めないプログラミング初心者ですふぇぇ(><)

先日のAtCoder Beginner Contest 128で、Python3を使って、AB2完を目指して解いていたところ、 B問題に79分取られるレベルで思考を停止させてしまったので、同じことを繰り返さないようにとメモを残しておくことにしました。
以下が私の提出したコードです(提出 #5649245)。


def irekae(array2):
    return [array2[1],array2[0],array2[2]]
n=int(input())
l=[input().split() for i in range(n)]
for i in range(len(l)):
    l[i].append(i+1)
    l[i][1]=int(l[i][1])
for i in range(len(l)):
    l[i]=irekae(l[i])
l.sort(reverse=True)
for i in range(len(l)):
    l[i]=irekae(l[i])
lx=sorted(l,key=lambda x: x[0])
for i in range(len(lx)):
    print(lx[i][2])

それでは自分の思考を追っていきたいと思います。

最初のアプローチ

さて、今回私が最初に考えたアプローチは、

都市名で辞書順ソートをする
同じ都市名のかたまりを作る
2でできた塊の中で点数をソートする

というものでした。途中でレストランに番号を紐付けなければいけないことに気づくことはできました。
とりあえず、まずは標準入力で受け取ります。


n=int(input())
l=[input().split() for i in range(n)]

ここに入力例1を入れると [['khabarovsk', '20'], ['moscow', '10'], ['kazan', '50'], ['kazan', '35'], ['moscow', '60'], ['khabarovsk', '40']] が返されます。

次に、レストランに番号を紐づけ、ついでにstr型になっていた点数をint型に直します。


for i in range(len(l)):
l[i].append(i+1)
l[i][1]=int(l[i][1])

ここで紐づけた番号は最後出力するときに必要になります。
これで、 [['khabarovsk', 20, 1], ['moscow', 10, 2], ['kazan', 50, 3], ['kazan', 35, 4], ['moscow', 60, 5], ['khabarovsk', 40, 6]] が返されます。

そして、最初のアプローチの手順1にあるように、都市名でソートすることにしました。
最初のリストは必要ないので、リストを書き換えてしまうsort()を使います。


l.sort()

これで [['kazan', 35, 4], ['kazan', 50, 3], ['khabarovsk', 20, 1], ['khabarovsk', 40, 6], ['moscow', 10, 2], ['moscow', 60, 5]] が返されました。

最初の躓き

しかし、あと一歩というところで私は躓いてしまいました。
というのも、手順2をどうやったら実装できるかがわからないのです。
とりあえず、点数を大きい順にソートすることだけはわかったので、それを楽にするために関数を定義しました。


def irekae(array2):
    return [array2[1],array2[0],array2[2]]

変数の名前がarray2なのは(いろいろあったので)無視してください。
この関数は、リストの0番目と1番目を入れ替えるはたらきをします。ここで読み込まれるリストは地名、点数、番号のリストなので要素数が3とはじめから決定されています。
ソートされるとき、「リストのリスト」の形では子リストの0番目の要素が参照される(これは間違いです、後述)と思いこんでいた私はこの関数をとりあえず定義しました。

本題に戻りましょう。
最初、私は [[['kazan', 35, 4], ['kazan', 50, 3]], [['khabarovsk', 20, 1], ['khabarovsk', 40, 6]], [['moscow', 10, 2], ['moscow', 60, 5]]] というように、地名ごとに塊を作ってその中で点数順にソートしようと考えていました。

ところが、地名ごとの塊を作る方法が思いつかなかったのです。

発想の転換

「配列をforで回してi番目とi+1番目の地名が違ったら分割…いやどうやったらそれができるんだ！？」みたいな状態で考え込んでリアルに1時間程度が経過した後、ふと私はひらめきました。

「先に点数をソートすればいいのでは？」

というわけで、先にあったように「リストのリスト」のソートは子リストの0番目の要素が参照されるという思い込みをもとに、irekae()で点数を0番目に持ってきました。そしてそれをソートします。点数が大きい順、つまり降順なのでsort()の引数reverseにTrueを指定します。


for i in range(len(l)):
    l[i]=irekae(l[i])
l.sort(reverse=True)

これで、[[60, 'moscow', 5], [50, 'kazan', 3], [40, 'khabarovsk', 6], [35, 'kazan', 4], [20, 'khabarovsk', 1], [10, 'moscow', 2]] が返ってきます。
これを今度は地名に基づいてソートすれば良いわけですから、私は先程の思い込みに基づいて地名を0番目に持ってこようとして、もう一度irekae()の後sort()を実行しました。
今度は地名を昇順にソートするのでreverseは無しです。


for i in range(len(l)):
    l[i]=irekae(l[i])
l.sort()

あとは問題のルールに基づき、l[i][2]を順番に出力すればACと思っていた私の元に返ってきたリストは、 [['kazan', 35, 4], ['kazan', 50, 3], ['khabarovsk', 20, 1], ['khabarovsk', 40, 6], ['moscow', 10, 2], ['moscow', 60, 5]] というものでした。

二度目の躓き

返ってきた子リストの1番目(つまり子リストの中央の要素)を見ると、せっかく降順にソートしたはずの点数が、昇順にソートし直されてしまっていました。ここで私は思い込みの間違いに気づきます。つまり、

誤:「リストのリスト」のソートは子リストの0番目の要素が参照される

正:「リストのリスト」のソートは最初に子リストの0番目の要素が参照されるが、0番目に重複があった場合は1番目以降も参照される

しっかり公式ドキュメントを読み込んだわけではないので断言はできないのですが、少なくとも私が見た限りではこの仕様が正しいです。
正解目前にして、またも痛い躓きです。

そして解決へ

ここで私がやるべきことは、「地名のみを参照してソートする」ことでした。「Python リスト」で検索して出てきたページを片っ端から舐めるように見ていくと、なんとか目的のものを見つけることができました。

Key関数 - ソート HOW TO — Python 3.7.3 ドキュメント

「よくある利用パターンはいくつかの要素から成る対象をインデクスのどれかをキーとしてソートすることです。」
ここで出された例は中身がリストじゃなくてタプルですが、やってることは全く同じです。というわけで、このように書きます。


lx=sorted(l,key=lambda x: x[0])
# 上記コード内ではなぜかsortedを使用しているがsortを使って以下のように書くことも可能
l.sort(key=lambda x: x[0])

lambdaは無名関数、つまり(厳密に言うと違いますが)使い捨てしやすい関数を定義します。そう何度も使わない、いちいちdefを使うほどでもない関数を定義するときなどに使うことができます。
ここでは変数xに子リスト(['khabarovsk', 20, 1]など)が代入され、返される値がx[0]、つまり0番目の要素です。その0番目の要素(ここでは地名)が参照されることで地名をkeyとしたソートが可能になります。
ここまできたらもう解決ですね。


lx=sorted(l,key=lambda x: x[0])
for i in range(len(lx)):
    print(lx[i][2])
# sortを使用した場合は以下のようになる
l.sort(key=lambda x: x[0])
for i in range(len(l)):
    print(l[i][2])

無事ACしました。

改善点

このコードにはいくつか改善すべき点がありました。それを指摘していきます。

sortとsortedを揃える

実は当初、keyパラメータはsortedでしか使えない！みたいな思い込みをしていたせいでこんなことになってしまっていますが、書いてあるとおり、普通にsortでも使うことは可能です。

配列の長さを計算する必要はない

このコードに何度も出てくるforループですが、回数指定のためのrange関数の中身は配列lの長さを指定しています。
ですが、実は最初に入力される値の数はnで指定されているため、len(l)は(len(lx)も当然)全てnに置き換えることができます。

irekae()は必要ない

irekae()はソートする対象を0番目に持ってくるための関数でしたが、sortおよびsortedに備わるkeyパラメータを使えば直接要素を指定してソートできるため、必要ありません。

工夫すればreverseする必要がない

これは解説のPDFを見て知ったのですが、点数が入力されるときに-1倍、つまり負の数にしておけば、点数の大小は正のときと逆になるため、reverseせずとも正しく点数を並べることができます。これ考えた人天才ですね。

以上を踏まえて、記号と文字の間にスペースを入れるなどして少しわかりやすく書くと、以下のようになります。


n = int(input())
l = [input().split() for i in range(n)]
for i in range(n):
    l[i].append(i+1)
    l[i][1] = -int(l[i][1])
l.sort(key = lambda x: x[1])
l.sort(key = lambda x: x[0])
for i in range(n):
    print(l[i][2])

圧倒的にきれいなコードになりました。

反省・まとめ

今回、私がB問題で79分も時間を取られてしまったのは、「地名の後に点数をソートする」という考えに囚われていたのが大きいです。先入観を排除して、柔軟に、いくつもの可能性を考えて問題にアプローチすることが大切だということを痛感しました。
次回以降のコンテストでもこの反省を活かせるよう努力するつもりです。

今回はこのブログでは初めてこんな真面目な記事を書きました。私自身決して詳しくないプログラミングに関する説明なので、間違いなども多くあるかもしれません。
そのようなものを見かけたら、遠慮なくTwitter等で指摘していただけると助かります。

Qman's Diary