はじめに

こんにちは。緑コーダーです。

競プロを始めたての頃は、ちょっとした入出力や変数の取り扱いに苦労しました。肝心のロジックまでなかなか手が届かず歯がゆい思いをしたものです。

しかし、多くは少しの知識や工夫で克服できるものでした。
そろそろ中級者の仲間入りと思いたい今日この頃。色変記事とかこつけて、過去の自分にチートシートを贈ります。

おことわり

実装上の工夫やパッケージの導入が主な内容です
アルゴリズムの解説はありません
より優れた実装があればやさしく教えてください

1. max(),min()

1.1 最大値（最小値）を更新する

ansを最大化する問題を考えます。頻出のシチュエーションですね。
変数testのほうが大きければansを更新します。
初心者時代の筆者は次のようにif文で実装していました。

before

ans = 0
if test > ans:
    ans = test

次のようにmax()関数を利用するとスマートに書けます。
最小値を求める問題であれば代わりにmin()を使います。

after

ans = 0
ans = max(ans, test)

上記の例ではわずか１行省略しただけですが、実際にはコード中の複数個所で最大値（最小値）が更新されることもよくあるため、地味ながら可読性に貢献してくれます。

1.2 最大値（最小値）を求める

当然ですが、複数の候補から最大値（最小値）を求めるときはif文よりも断然max(), min()が便利です。

例： ABC395 B

[概要]
整数$N$が与えられます。
次のような$N \times N$の模様を描画してください。
$N=5$の例：
#####
#...#
#.#.#
#...#
#####

要するに一番外の枠は#、内側の枠は.、さらに内側の枠は#という要領で年輪のような模様を構築する問題です。グリッド上の点$(i,j)$に対応する記号は、上下左右の端からの距離の最小値の偶奇で判定できます。

ifで分岐するのはさすがに現実的ではありませんね。

解答の方針

for i in range(N):
    for j in range(N):
        d = min(i, j, N-1-i, N-1-j)
        S[i][j] = "#" if d%2==0 else "."

2. 座標の取り扱い

2-1. 近接する座標を探索する場合

[例題]
グリッド上の点$(i,j)$について、上下左右の座標に対して操作を行ってください。

よくある設定です。肝心の『操作』の部分は問題によってさまざまで、

対象のマスに色を塗る
対象のマスに移動する
対象のマスを連結させる...etc

などのバリエーションがあります。今回は別途用意されたcalc()関数を作用させることにしますが、適宜読み替えてください。
初心者時代には次のように実装していました。

before

calc(i+1, j)
calc(i-1, j)
calc(i, j+1)
calc(i, j-1)

いかがでしょうか。
$(i,j)$に対して上下左右の４座標をすべて列挙しています。間違いではありませんが、問題設定の影響を如実に受けますし、隣り合う座標の計算でバグを埋め込みやすい、やや窮屈な書き方です。

次のように、座標の移動先を管理する変数を用意すると見通しが良くなります。

after

di = [1, 0, -1, 0]
dj = [0, 1, 0, -1]

for d in range(len(di)):
    ni, nj = i+di[d], j+dj[d]
    calc(ni, nj)

一番のメリットは移動先に関する情報をまとめておけることです。これにより、変則的な問題にも$di$と$dj$を書き換えるだけで柔軟に対応できます。移動先を間違えた場合も$di$と$dj$を確認するだけでOKです。

例１：ABC377 C チェスのナイトの動き

チェスのナイトの動き

di = [1, 2, 2, 1, -1, -2, -2, -1]
dj = [2, 1, -1, -2, -2, -1, 1, 2]

for d in range(len(di)):
    ni, nj = i+di[d], j+dj[d]
    calc(ni, nj)

例２：ABC413 G 上下左右斜めの８方向

八方

di = [0, 1, 1, 1, 0, -1, -1, -1]
dj = [1, 1, 0, -1, -1, -1, 0, 1]

for d in range(len(di)):
    ni, nj = i+di[d], j+dj[d]
    calc(ni, nj)

forループの中身を全く同じ書き方にできますね。

2-2. 方向が文字で指定される場合

例：ABC421 D

[設定の抜粋]
文字$S$に応じて座標$(i,j)$から次のように移動します。

$S$がUのとき$(i-1,j)$に移動

$S$がDのとき$(i+1,j)$に移動

$S$がRのとき$(i,j+1)$に移動

$S$がLのとき$(i,j-1)$に移動

素直にif文で実装できますが、いかにも面倒です。

before

if S == "U":
    calc(i-1, j)
if S == "D":
    calc(i+1, j)
if S == "R":
    calc(i, j+1)
if S == "L":
    calc(i, j-1)

辞書で移動先をまとめて扱うと簡潔です。

after

D = {"U":(-1,0), "D":(1,0), "R":(0,1), "L":(0,-1)}
di, dj = D[S]
calc(i+di, j+dj)

3. 入力の受け取り方

3.1 複数行の入力

$N$人のお客さんがレストランを訪れます。
$人i$の入店時刻は$A_i$、退店時刻は$B_i$です。
[入力例]
$N$
$A_1$ $B_1$
...
$A_N$ $B_N$

非常によくある問題設定です。
このようなお客さん情報を変数$C$に格納してみましょう。

before

N = int(input())
C = []
for _ in range(N):
    A, B = map(int, input().split())
    C.append((A,B))

いちいちループを書くのは面倒ですね。
内包表記 を利用すると、配列への格納が１行で済みます。

after

N = int(input())
C = [tuple(map(int, input().split())) for _ in range(N)]

内包表記は分かりにくいという声を聞きます。否定はしません。
コードの簡潔さと可読性はトレードオフになる部分があり、無理に真似ろと主張する意図はありません。ただ、慣れれば便利なので紹介します。

3.2 入力値に対して一定の操作を施す

グリッド上に$N$個の点があります。
点$i$の座標は$(X_i,Y_i)$で与えられます。
[入力例]
$N$
$X_1$ $Y_1$
...
$X_N$ $Y_N$

問題によっては入力値が1-indexで与えられるので、あらかじめ0-indexに直す必要があります。
１を引く操作を記述するため、結局次のようにループで実装することになるのでしょうか。

before

N = int(input())
C = []
for _ in range(N):
    X, Y = map(int, input().split())
    C.append((X-1,Y-1))

よく見ると$X$、$Y$ともに、整数に直して１引くだけの処理です。
このように一定の操作を施すだけならlambda関数が便利です。map()の第一引数にlambdaを指定します。

after

N = int(input())
C = [tuple(map(lambda x: int(x)-1, input().split()))]

3.3 辞書の初期値

$N$頂点$M$辺のグラフがあります。
$i$番目の辺は$u_i$と$v_i$を結ぶ無向辺です。
[入力例]
$N$ $M$
$u_1$ $v_i$
...
$u_N$ $v_N$

グラフの問題も頻出ですね。
辞書$G$にこのような辺の情報を格納するとします。素直に実装すると次のようになります。

before1

N, M = map(int, input().split())
G = {}
for _ in range(M):
    u, v = map(lambda x: int(x)-1, input().split())
    if u in G:
        G[u].append(v)
    else:
        G[u] = [v]
    if v in G:
        G[v].append(u)
    else:
        G[v] = [u]

辞書オブジェクトに存在しないキーでアクセスするとエラーになってしまうため、キーに$u$や$v$が存在するかチェックしています。
面倒ですね。$u$と$v$に対して別々に判定を実施する必要があり、コードが冗長になっています。
次のように工夫すればどうでしょうか。

before2

N, M = map(int, input().split())
G = {i:[] for i in range(N)}
for _ in range(M):
    u, v = map(lambda x: int(x)-1, input().split())
    G[u].append(v)
    G[v].append(u)

内包表記を使って$G$を宣言しました。中身はG = {0:[], 1:[], ..., N-1:[]}このようになっています。$0$から$N-1$までの整数をキーとしているため、面倒だったif文を削除できました。めでたしめでたし。

……というのは罠です。$N$が$O(10^5)$程度なら問題なく機能すると思いますが、$O(10^{10})$など非常に大きいとタイムアウトやメモリ超過になります。リストもそうですが、要素が多すぎると制限内で扱いきれないのですね。
また、文字列を頂点としたグラフの場合に応用が利きません。

上記２つの問題を、 defaultdict なら解決できます。

after

from collections import defaultdict
N, M = map(int, input().split())
G = defaultdict(list)
for _ in range(M):
    u, v = map(lambda x: int(x)-1, input().split())
    G[u].append(v)
    G[v].append(u)

はじめてキーにアクセスしたとき、自動的に要素が生成されます。つまりキーの存在判定が不要です。
また、アクセスしなかったキーに対しては要素が作られないため、実行時間やメモリにやさしいです。

初期値は引数のコンストラクタを使って生成されます。リストだけでなく、問題に合わせて設定できます。

いろいろな初期値

from collections import defaultdict
# 初期値０
G = defaultdict(int) # int() = 0 になるため

# 初期値10**10
G = defaultdict(lambda: 10**10)

# 初期値inf
inf = float("inf")
G = defaultdict(lambda: inf)

3.4 可変長の入力

$Q$個のクエリが与えられます。クエリは次の３種類のうちいずれかです。

$1$ $X$ $Y$

$2$ $X$

$3$ $Y$

この手の問題が非常に苦手でした。とりあえず入力をリストで受け取るとして、３番目の値があるかどうかの判定をして、２番目の値は$X$だったり$Y$だったり......なんてやってられませんね。せっかくの競プロ、問題に集中したいです。

次のように アンパック を利用しましょう。

after

Q = int(input())
for i in range(Q):
    query = list(map(int, input().split()))
    typ = query[0]
    if typ == 1:
        _, X, Y = query # こういうのがアンパック
    elif typ == 2:
        _, X = query
    elif typ == 3:
        _, Y = query

自動的に左辺の変数たちへ右辺の中身を分配してくれます。リストの先頭の値はifの分岐だけに利用されるので、if節の中ではもう使わないよという意味を込めて_にしています。

右辺の要素数と左辺の変数の数をそろえる必要がある点には注意してください。

3.5 複数個の変数に格納する

上の3.4の派生です。
グラフ問題でのbfs/dfsなど、キューで座標を管理するケースを考えます。

before

while len(q)>0:
    pos = q.popleft()
    i = pos[0]
    j = pos[1]

初心者時代にはこのように、キューから取り出した変数をいちいち$i,j$に分配していました。キューに格納する変数の数がふえると厄介ですね。

ここでもアンパックの記法が活躍してくれます。

after

while len(q)>0:
    i,j = q.popleft()

記述量を減らせたうえ、一瞬使われるだけの変数posを除去できました。

おわりに

上級者にとっては当たり前の内容だったかもしれません。
しかし、実装を省略したり楽できたりすることは、案外気づかないことも多々あるかと思います。この記事がだれかの助けになれば幸いです。

それでは皆さん良き競プロライフを！