Python
AtCoder
競技プログラミング
VisualStudioCode

Pythonで競プロやるときによく書くコードをまとめてみた

はじめに

最近,よくPythonを使って競技プログラミングに参加しています.そこで競プロの問題を解くときに書くことの多い単純な処理について,Python3のコードを基本的なところからまとめてみます.文字列操作やデータ操作に慣れると,本質的なアルゴリズムに専念できてより楽しくなると思います.競プロやってみたいなと思っている人の参考になればうれしいです.(もちろん競プロ以外でも参考になると思います)

まずは基本

まずは基本的な入出力からです.ここで躓くと面白くないのでしっかりと押さえておきたいところです.

基本的な入出力

最近自分がはじめたAtCoderや2018年のGoogle Code Jamは入力データを標準入力で受け取り,回答を標準出力に出して結果を判定してもらいます.Pythonでは入力はstdin.readline(),出力はprint()を使うとうまくいきます.
入力データの末尾に改行文字が入るのが邪魔になることもあるのでrstrip()を使うとよいです.

spam.py
from sys import stdin

a = stdin.readline().rstrip()
print(a.upper())    # 読み込んだ文字を大文字に変換

ローカルでテストするときは入力データをファイルに書いておいて,リダイレクトして読み込むと簡単です.

sample.txt
spam!!
実行
$ python spam.py < sample.txt
SPAM!!

またインタラクティブな問題の場合は出力時にflushする必要がありますが,print関数のオプションで指定できます.

print('spam', flush=True)

入力値の簡単なデータ操作

標準入力から受け取った文字列はそのままでなく,何らかの変換をすることが多いです.それらについてまとめていきます.

数値への変換

これはよくありますね.int()を使います.

n = int(stdin.readline().rstrip())

空白で分割

これもよくあります.split()の結果はリストで戻ってきますが,問題によってはアンパック代入を使ったほうが分かりやすい場合も多いので覚えておくとよいです.区切り文字はsep引数で指定ができますが,競プロでよく使われる空白区切りの場合は指定なしでOKです.

spam.py
l = stdin.readline().rstrip().split()
A, B, C = stdin.readline().rstrip().split() # 問題文に出てくるA,B,Cが空白区切りで1行で入力される場合はアンパック代入が便利
print(l)
print(A, B, C)
sample.txt
1 2 3
AAA BBB CCC
実行
$ python spam.py < sample.txt
['1', '2', '3']
AAA BBB CCC

空白で分割しつつ数値に変換

上のふたつの合わせ技ですが,リスト内包表記を使うと1行で書けます.リスト内包表記を覚えるとPythonが楽しくなると思います.

spam.py
A, B, C = [int(x) for x in stdin.readline().rstrip().split()]
print(A + B - C)
sample.txt
1 2 3
実行
$ python spam.py < sample.txt
0

リスト内包表記を使ったあれこれ

先ほど出てきたリスト内包表記を使ったデータ操作についてまとめていきます.

特定の条件のデータを数える

リストの中からある条件を満たす要素の個数をカウントするときもリスト内包表記が使えます.リスト内包表記のifを使って条件を満たすリストを作り,len()でそのリストの長さを求めます.

spam.py
data = range(1, 10)
count = len([x for x in data if x % 3 == 0]) # 3の倍数をカウント
print(count)
実行
$ python spam.py
3

for文で回しながら条件チェックを行いcount変数をインクリメントするより簡潔に書けます.

複数行をまとめて読み込む

複数行のデータをまとめて読み込むときもリスト内包表記を利用できます.ここでは入力データの1行目にデータの行数,2行目以降に空白区切りのデータでマトリクス状のデータが入力される例をあげます.

spam.py
n = int(stdin.readline().rstrip())
data = [stdin.readline().rstrip().split() for _ in range(n)]
print(data)
sample.txt
3
1 2 3
4 5 6
7 8 9
実行
$ python spam.py < sample.txt
[['1', '2', '3'], ['4', '5', '6'], ['7', '8', '9']]

また,リストのオブジェクトを上記のようにそのままprint関数に渡すと1行で表示されますが,print関数の呼び出しで少し工夫するとマトリクス状のイメージそのままに表示されて確認しやすいです.

spam.py
print(*data, sep='\n')
実行
$ python spam.py
['1', '2', '3']
['4', '5', '6']
['7', '8', '9']

リストのオブジェクトに*をつけて引数に渡すと分解して渡されるのを利用しています.またsepオプションに改行文字を指定することで,分割して渡したデータを改行しながら表示しています.

二重のリストの行列を入れ替える

マトリクス状のデータを横方向(行)だけでなく縦方向(列)にも操作したいときには,行列の転置を行うと操作しやすいです.また競プロでは外部ライブラリを使えないこともあるので,numpyを使わずにzipを使ってリスト内包表記で解決します.

spam.py
data = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
print(data)

transpose = [x for x in zip(*data)]
print(transpose)
実行
$ python spam.py
[[1, 2, 3], [4, 5, 6], [7, 8, 9]]
[(1, 4, 7), (2, 5, 8), (3, 6, 9)]

これを利用すると列ごとの和を求めたりするのも簡単です.

spam.py
data = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
print([sum(x) for x in zip(*data)])
実行
$ python spam.py
[12, 15, 18]

二重のリストをフラットにする

マトリクス状のデータを1行の単純なリストに平坦化するときにもリスト内包表記が使えます.

spam.py
data = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
print(data)

flat = [flatten for inner in data for flatten in inner]
print(flat)
実行
$ python spam.py
[[1, 2, 3], [4, 5, 6], [7, 8, 9]]
[1, 2, 3, 4, 5, 6, 7, 8, 9]

ちょっと変わった他の方法としてsum()を使ってフラットなリストを作ることもできます.第二引数で合計する最初のデータを指定できるのですが,ここで空リストを渡しておくと+演算子でリストの連結が行われるようです.

spam.py
data = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
print(sum(data, []))
実行
$ python spam.py
[1, 2, 3, 4, 5, 6, 7, 8, 9]

その他のリストまわりの操作

リスト内包表記の他にも,覚えておくとよいリスト関連のデータ操作があるのでまとめます.

連結と繰り返し

基本的なものですが押さえておきたい操作です.

spam.py
a = [1, 2, 3] + [4, 5]
print(a)

b = [1, 2, 3] * 3
print(b)
実行
$ python spam.py
[1, 2, 3, 4, 5]
[1, 2, 3, 1, 2, 3, 1, 2, 3]

要素の存在確認

こちらも基本的な操作ですが,知らないとわざわざfor文をまわしてしまうことになるので押さえておきたいです.

spam.py
a = [1, 2, 3, 4, 5]
print(1 in a)
print(6 in a)
実行
$ python spam.py
True
False

削除

リストから要素を削除する方法はいくつかあります.状況に応じて適切な方法を使えるようにしたいです.

spam.py
a = [1, 2, 3, 4, 5]
del a[1]        # インデックス指定で削除したい場合
print(a)

b = [1, 2, 3, 4, 5]
del b[1:3]      # スライスで部分リストを指定して削除も可能
print(b)

c = [1, 2, 3, 4, 5]
x = c.pop(1)    # popでもインデックス指定で削除可能
print(c, x)

d = [1, 2, 3, 4, 5]
d.remove(3)     # オブジェクトを指定して削除したい場合
print(d)

実行
$ python spam.py
[1, 3, 4, 5]
[1, 4, 5]
[1, 3, 4, 5] 2
[1, 2, 4, 5]

ソート

もとのリスト自体を書き換えて破壊的にソートするsort()メソッドと,ソートした結果のリストを別途生成して返すsorted()関数があります.状況に応じて使い分けるといいと思います.

spam.py
a = [2, 5, 1, 4, 3]
b = sorted(a)

print(a)    # a自体は変更なし
print(b)

a.sort()
print(a)    # a自体がソートされている
実行
$ python spam.py
[2, 5, 1, 4, 3]
[1, 2, 3, 4, 5]
[1, 2, 3, 4, 5]

降順にソートしたい場合は,reverseオプションをTrueに指定します.

spam.py
a = [2, 5, 1, 4, 3]
b = sorted(a, reverse=True)
print(b)
実行
$ python spam.py
[5, 4, 3, 2, 1]

keyオプションに比較を行う前にリストの各要素に対して呼び出される関数を指定することで,柔軟なソートを行うことができます.

spam.py
a = ['e', 'B', 'd', 'C', 'a']
print(sorted(a))                            # 普通にソートすると大文字,小文字それぞれでソートされます.
print(sorted(a, key=lambda x: x.upper()))   # 比較時に大文字に変換してソートするので,結果大文字・小文字あわせてソートされます.
実行
$ python spam.py
['B', 'C', 'a', 'd', 'e']
['a', 'B', 'C', 'd', 'e']

もう一つkeyを指定する例として,多重キーでのソートをあげます.内部のタプルの1要素目(アルファベット),2要素目(数字),0要素目(数値)の順に比較してソートを行います.

spam.py
a = [(1, 'One', '1'), (1, 'One', '01'),
     (2, 'Two', '2'), (2, 'Two', '02'),
     (3, 'Three', '3'), (3, 'Three', '03')]
print(sorted(a, key=lambda x: (x[1], x[2], x[0])))
実行
$ python spam.py
[(1, 'One', '01'), (1, 'One', '1'), (3, 'Three', '03'), (3, 'Three', '3'), (2, 'Two', '02'), (2, 'Two', '2')]

マイナス値のインデックス

Pythonのリストは,インデックスを指定してアクセスするときにマイナス値を使うことができます.マイナスを指定すると末尾から数えられ,-1が末尾になります.通常はリストのLengthを求めて末尾を指定することが多いですが,マイナス値のインデックスを使うとスマートに書くことができます.

spam.py
a = [1, 2, 3, 4, 5]
print(a[len(a) - 1])    # よくある末尾の指定
print(a[-1])            # マイナス値のインデックスを使った末尾の指定
print(a[-2])            # 末尾から2番目
実行
$ python spam.py
5
5
4

部分リストの取り出し(スライス)

Pythonでリストの中の一部分を取り出すときにはスライスを使うと簡単に取得することができます.[始点:終点]の形で指定します.(終点自身は含まれません.始点以上~終点未満です.)
始点と終点は省略も可能でその場合は先頭から,あるいは末尾までの取り出しになります.

spam.py
a = [1, 2, 3, 4, 5, 6, 7, 8, 9]
print(a[3:7])
print(a[3:])    # 終点の省略
print(a[:7])    # 始点の省略
実行
$ python spam.py
[4, 5, 6, 7]
[4, 5, 6, 7, 8, 9]
[1, 2, 3, 4, 5, 6, 7]

またステップを指定することもできます.

spam.py
a = [1, 2, 3, 4, 5, 6, 7, 8, 9]
print(a[2:8:2])     # 2番目以上8番目未満を2step(1つ飛ばし)で取得
実行
$ python spam.py
[3, 5, 7]

ステップにはマイナス値を指定することもでき,後ろ側から順に取り出します.その場合は始点のほうに大きいインデックスを指定して使います.

spam.py
a = [1, 2, 3, 4, 5, 6, 7, 8, 9]
print(a[6:1:-2])    # 6番目から2番目まで(1番目は含まれない)を2stepで減らしながら取得
実行
$ python spam.py
[7, 5, 3]

シャローコピーとディープコピー

コピーするときにはcopyモジュールが使用できますが,リストの単純なコピーには全スライスを使って簡単に書くことができます.

spam.py
a = [1, 2, 3, 4, 5]
b = a[:]    # 終点始点を指定してすべてをスライスで取得(結果コピーと同じ)
b[0] *= 10  # bはaのコピーなのでaには影響がありません.

print(a)
print(b)
実行
$ python spam.py
[1, 2, 3, 4, 5]
[10, 2, 3, 4, 5]

ただし多重リストをこの方法でコピーした場合,単純なシャローコピーになるので思わぬ副作用が出てBugになってしまうことがあります.

spam.py
a = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
b = a[:]        # シャローコピー
b[0].append(0)  # b[0](内部のリスト)はaと共通のものを参照している

print(a)        # aにも0が追加されている.(期待しない結果!)
print(b)
実行
$ python spam.py
[[1, 2, 3, 0], [4, 5, 6], [7, 8, 9]]
[[1, 2, 3, 0], [4, 5, 6], [7, 8, 9]]

これを防ぐにはcopy.deepcopyを使ってディープコピーすると解決します.

spam.py
a = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
b = deepcopy(a)     # ディープコピー
b[0].append(0)      # b[0](内部のリスト)はaとは別のオブジェクト

print(a)            # aに影響はない
print(b)
実行
$ python spam.py
[[1, 2, 3], [4, 5, 6], [7, 8, 9]]
[[1, 2, 3, 0], [4, 5, 6], [7, 8, 9]]

和集合と積集合

2つのリストで少なくとも片方に含まれているものを集めたものが和集合,両方ともに入っているものを集めたものが積集合です.これらはセットを使うと簡単に求めることができます.リストとセットの変換はlist()set()を使って行います.

spam.py
a = [2, 4, 6, 8]
b = [3, 6, 9]

print(list(set(a) | set(b)))  # 和集合
print(list(set(a) & set(b)))  # 積集合
実行
$ python spam.py
[2, 3, 4, 6, 8, 9]
[6]

文字列操作

リストの操作と同じく,文字列の操作もよく使うので覚えておくとよいです.

リストと同じ操作

連結と繰り返し,存在確認などリストと同じ操作が文字列でも行えます.

spam.py
a = 'ABC' + 'DEF'
print(a)

b = 'ABC' * 3
print(b)

c = 'ABCDEF'
print('A' in c)
print('a' in c)
実行
$ python spam.py
ABCDEF
ABCABCABC
True
False

置換

基本的な操作ですが文字列の置換についてです.また大文字・小文字の変換も簡単に行えます.文字列操作をするときに大文字・小文字の区別をしないときには,あらかじめどちらかに変換しておくと便利な時があります.

spam.py
a = 'abc(de)fg'
print(a.replace('(', '[').replace(')', ']'))

b = 'aBCdEfg'
print(b.upper())
print(b.lower())
実行
$ python spam.py
abc[de]fg
ABCDEFG
abcdefg

文字列の反転

競プロでは文字列を末尾から見ていくとうまく解ける問題が出てくることがあり,文字列の反転のやり方を覚えておくと便利です.Pythonの文字列はリストと同じくスライスが使えるので,文字列の反転に利用できます.

spam.py
s = 'ABCDEFG'
print(s[::-1])
実行
$ python spam.py
GFEDCBA

反転に限らず部分文字列の取り出しにもスライスが使えるので覚えておくとよいです.

文字列定数

stringモジュールにはいくつかの文字列定数があります.競プロのように限られた時間でコードを書く場合は特に,これらの定数を自前で書かずに用意されているものを利用することで,手早くコードを書くことができます.

spam.py
print(string.ascii_lowercase)
print(string.ascii_uppercase)
print(string.ascii_letters)
print(string.digits)
print(string.hexdigits)
実行
$ python spam.py
abcdefghijklmnopqrstuvwxyz
ABCDEFGHIJKLMNOPQRSTUVWXYZ
abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ
0123456789
0123456789abcdefABCDEF

正規表現

文字列操作といえばということで正規表現についても簡単な例をまとめておきます.reモジュールを使用します.

spam.py
import re

text = '<h1 style="width: 100px; height: 200px;">'
result = re.search(r'width: (\d*)px; height: (\d*)px;', text)

if result:
    print(result.group(0))
    print(result.group(1))
    print(result.group(2))
    print(result.groups())
else:
    print('no match')
実行
$ python spam.py
width: 100px; height: 200px;
100
200
('100', '200')

itertoolsを使った操作

itertoolsモジュールの中には,自前で算出するにはちょっと面倒そうな操作がいくつかあります.覚えておくと役に立つと思います.

累積和

競プロでは累積和を使うことで効率よく解くことができる問題をよく見かけます.自前で求めるのも難しくありませんが,よく使うのでaccumulateを知っておくと便利です.

spam.py
from itertools import accumulate

a = list(range(1, 11))
b = list(accumulate(a))     # itertoolsの戻り値はイテレータとなっているので必要に応じてlist化します.
print(a)
print(b)
実行
$ python spam.py
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
[1, 3, 6, 10, 15, 21, 28, 36, 45, 55]

dropwhileとtakewhile

dropwhileは先頭から与えられた条件が真である限り除外して,それ以降の要素を戻します.takewhileはその逆で先頭から与えられた条件が真であるところまでの要素を戻します.

spam.py
from itertools import dropwhile, takewhile

a = [3, 6, 1, 7, 2, 5]

b = dropwhile(lambda x: x != 1, a)  # 1が出るまでを除外する
print(list(b))

c = takewhile(lambda x: x != 1, a)  # 1が出るまでを取り出す
print(list(c))
実行
$ python spam.py
[1, 7, 2, 5]
[3, 6]

groupby

連続する要素をグループ化します.ただし同じ要素でも連続していないものはグループ化しません.

spam.py
from itertools import groupby

a = [1, 1, 2, 3, 3, 3, 1, 2, 2]

for key, value in groupby(a):
    print(key, list(value))
実行
$ python spam.py
1 [1, 1]
2 [2]
3 [3, 3, 3]
1 [1]
2 [2, 2]

またsort等と同様に,keyオプションに比較を行う前にリストの各要素に対して呼び出される関数を指定することができます.次の例はkeyに2の余りを算出するラムダを指定して,奇数・偶数ごとにグループ化しています.

spam.py
from itertools import groupby

a = [1, 3, 2, 4, 3, 1, 1, 2, 4]

for key, value in groupby(a, key=lambda x: x % 2):
    print(key, list(value))
実行
$ python spam.py
1 [1, 3]
0 [2, 4]
1 [3, 1, 1]
0 [2, 4]

連続していないものもまとめてグループ化したい場合は,あらかじめソートしてからgroupbyするとよいです.

spam.py
from itertools import groupby

a = [1, 1, 2, 3, 3, 3, 1, 2, 2]
b = sorted(a)

for key, value in groupby(b):
    print(key, list(value))
実行
$ python spam.py
1 [1, 1, 1]
2 [2, 2, 2]
3 [3, 3, 3]

その他

これまでにまとめてきたもの以外で思いつくコードについてまとめます.

for文のelse

Pythonのfor文にはelseをつけることができます.elseの中はfor文から抜ける時に実行されます.ただし途中でbreakしたときには実行されません.フラグを使わずにスマートに書くことができます.

spam.py
for i in range(5):
    if i == 5:
        break;      # このforではbreakされることがない
else:
    print('@@@')

for i in range(10):
    if i == 5:
        break;      # このforでは途中でbreakされる
else:
    print('###')
実行
$ python spam.py
@@@

2進数や16進数

10進数と2進数または16進数の文字列への変換はbin()hex()関数を使用します.10進数への変換はint()関数の第2引数に基数を渡して行います.

spam.py
print(bin(255))                 # 10進数 -> 2進数
print(hex(255))                 # 10進数 -> 16進数

print(int('0b11111111', 2))     # 2進数 -> 10進数
print(int('0xff', 16))          # 16進数 -> 10進数
実行
$ python spam.py
0b11111111
0xff
255
255

順列数と組み合わせ数

全順列,または組み合わせパターンを求めて総当たり,,,というと競プロでは望み薄なケースが多いと思いますが,順列数や組み合わせ数を求めて解く問題はあります.これにはmath.factorial()を使って計算することができます.

spam.py
from math import factorial

def permutations_count(n, r):
    ''' 順列 '''
    return factorial(n) // factorial(n - r)

def combinations_count(n, r):
    ''' 組み合わせ '''
    return factorial(n) // (factorial(n - r) * factorial(r))

print(permutations_count(5, 2))
print(combinations_count(5, 2))
実行
$ python spam.py
20
10

メモ化

メモ化とは一度計算した結果を覚えておいて,次に同じ引数で呼ばれた場合に計算をせずに覚えておいた値を戻します.@lru_cacheデコレータを使うことで簡単にメモ化することができます.知っておくと役に立つことがあるかもです.

spam.py
from functools import lru_cache

@lru_cache(maxsize=None)
def fibo(n):
    if n <= 0:
        return 0
    elif n == 1:
        return 1
    else:
        return fibo(n - 1) + fibo(n - 2)

print(fibo(35))
実行
$ python spam.py
9227465

stdinをファイルに置き換える

通常は冒頭のほうで説明したように,入力データをファイルに保存しておいてリダイレクトでstdinにデータを流し込んでいますが,そのように実行ができない場合はstdinをファイルに置き換えたい場合があります.
例えば自分はPythonのコードを書く場合はVisual Studio Codeを使っているのですが,Debuggerを使おうと思うとリダイレクトして実行することができません.その場合はファイルをopenした戻り値でstdinを上書きすることで解決します.

from sys import stdin
stdin = open('sample.txt')          # 冒頭にこの1行を追加

print(stdin.readline().rstrip())    # sample.txtから1行読み込まれる

おわりに

こういった単純なデータ操作などに時間と労力を使わずにさっとかけるようになると,競技プログラミングの本質的なアルゴリズムを考えるところにより多くの力を割くことができると思います.競技プログラミングをやってみたいと思っている人,また本質的でないところのコードを書くのに時間がかかってしまっている人などの参考になれば幸いです.