PSM

python する man

pythonで競プロやるときのメモ

python競技プログラミングの問題を解いていて,テクニックをまとめました.
ちょくちょく使うのに覚えてなくて,その度になんだっけって調べるやつです.

pythonって2だったり3だったり紛らわしいし,調べてみてもコピペで動いてくれるスクリプトが載っている記事を見つけられなかったので,自分の検索能力のなさを嘆きつつ苦しんだ結果をメモ書き程度に残しておきます.
(環境はpython3.6.0を想定しています)

入力

1行に1つの入力を受け取る

28

とか

test

とかを受け取りたい時です.
たぶんどの言語でもそうなんでしょうけど,pythonでも入力関数を一回呼び出すごとに1行分の入力を読み込んでくれるようです.
というわけで,

a = input()

とすればaに入力を代入できます.

ただしこれだとintだろうとfloatだろうとstrだろうと何でもstr型で読み込んでしまうので,型変換までついでにやっちゃうなら

a = int(input())

のように書けます.

1行に複数の入力を受け取る

a b c d e

1 2 3 4 5

とかを受け取りたい時です.

文字列で受け取るなら

a = input().split()

とすると,半角スペースのところで区切ってリストにして受け取ってくれます.

型を指定して読み込むときには,変数が1つのときの場合のようには書けず,map関数を使う必要があります.
intで受け取りたいなら

a = list(map(int, input().split()))

となります.
注意点としては,python3ではmap()を使って読み込むとmap objectというものになってしまうので,list()で変換する必要があります.

また,受け取りたい入力が数個程度で,その数がわかっているならこんな風にもできました
たとえば

3 4

に対して

a,b = list(map(int, input().split()))

とすればaに3が,bに4がint型で代入されます.


複数行に複数の入力を受け取る

3
1 2 3
4 5 6
7 8 9

競プロでよくある,みたいなのを受け取りたい時です.

input()を一回呼び出すごとに1行分の入力を受け取ってくれるので,for文などで行数分呼び出すか,whileなどを使えば良いようです.
というわけで,

n = int(input())
a = []
for i in range(n):
    a.append(list(map(int, input().split())))

最初に空のリストを用意して,appendでどんどん突っ込んでいくという方針です.
numpyを使うならここでついでにnp.arrayにしておいても良いかもしれません.


出力

基本

python3では標準出力printが関数となっていました(python2系との違い).

print("output")

のように書きます.
勘違いしていたのですが,ここではint型も勝手に型変換して出力してくれるようなので,例えば

print(1234)

とint型を直接書いても,勝手に変換して"1234"と出力してくれます.
複数の引数をとることもできて,

n = 2
print("春が",n, "階から降ってきた")

とすれば

春が 2 階から降ってきた

リストを空白などの文字区切りで出力する

join()を用いると,リストの要素を任意の文字を糊にしてくっつけた文字列に変換することができます.
["君","羊","青"]を"と"でくっつくて"君と羊と青"にするイメージです.

s = ["君","羊","青"]
joined_s = "と".join(map(str, s))
print(joined_s)

とすれば

君と羊と青

"と"の部分を半角スペース" "などにすれば半角スペース区切りで出力できます.


リスト操作

文字列を1文字ごとのリストに変換する

list()をかませるだけ

print(list("letters"))
['l', 'e', 't', 't', 'e', 'r', 's']

リストを1つの文字列に変換する

ほかのとこでも書いた気がする.joinを用いるのが定石.

l = ['l', 'e', 't', 't', 'e', 'r', 's']
print("".join(l))
letters

リスト内の要素の個数を数える

collectionsを使うのが早いし簡単らしい.

a = [1,2,3,5,2,1,3,4]
import collections
c = collections.Counter(a)

数えた結果は以下のように呼び出せる.

print(c)
print(c[2])
print(c.keys())
print(c.values())
Counter({1: 2, 2: 2, 3: 2, 5: 1, 4: 1})
2
dict_keys([1, 2, 3, 5, 4])
dict_values([2, 2, 2, 1, 1])

英小文字,英大文字のリストの作り方

アルファベットって何文字あるんだか,たまにわからなくなる.

lowl = [chr(i) for i in range(97, 97+26)]
uppl = [chr(i) for i in range(65, 65+26)]
['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z'] 
 ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z']

二次元リストをフラットにする

itertoolsを用いると早い.

l = [[1,2,3], [4,5,6]]
from itertools import chain
flatten_l = list(chain.from_iterable(l))
[1, 2, 3, 4, 5, 6]

deque

listと比べて良いところ.
listは,最後尾の要素をいじるappendとかの操作は早いんですが,先頭要素を加えたり削除したりしようとすると途端に遅くなります.
この点,dequeは双方向連結リストなので,先頭要素を操作するのも高速で行うことができるようになっています.

準備

collectionsをimportし,aという名前で dequeを準備します

import collections.deque
a = deque([1,2,3,4,5])

最後尾に要素を加えたい/削除したい

listと同じで,append/popを使います

#最後尾に要素を加える
a.append(0) #要素を加える
print(a)

#最後尾から要素を取り除く
print(a.pop()) #要素そのものも取り出せる
print(a) 

先頭に要素を加えたい/削除したい

pythonは左から右に流れる言語なので(逆は存在するのだろうか?)
先頭ということは即ち最も左側であるということになります.
よってappend/popにleft(左)を付け加えたappendleft/popleftを使います.

a.appendleft(9)
print(a)

print(a.popleft())
print(a)

map関数とlambda式

map関数とは

競技プログラミングをやっていて,標準入力を受け取るときに何も考えずに

a = list(map(int, input().split()))

とか書いたりしてますが,実はあまりよくわかっていませんでした.
map()というのは,1つ目の引数で指定した名前の関数を,2つ目の引数で指定したリストの要素一つ一つに当てはめていってくれる関数のようです.
つまり,ここではinput().split()で受け取られた入力のリスト(str型)を,一つづつ全てint型に変換してくれるという処理を行なっているわけでした.

ここではintという関数を使っていますが,この部分には自分で定義した関数を用いることができます.
例えば引数の2乗を返す関数fを定義して,

def f(a):
    return a**2

list(map(f, [1,2,3,4]))

とすれば

[1, 4, 9, 16]

とやってくれます.便利ですね.
注意点ですが,map関数はそのままでは結果をmap objectという形式で返すようになっているため,printなどがしたい場合はlistをかませる必要があります.

lambda式

ところで,上の式では自作関数fをいちいち定義していましたが,これをもっとスマートにやる書き方があります.
それがlambda式です.
正確に言うとlambda式は,無名関数を定義する記法の1つであり,関数の定義文を式として扱うことができるのです.

lambda 変数:返り値の定義

のように書きます.

lambda a: a>0

すると,

<function __main__.<lambda>>

ちゃんと関数として定義されているようです.

この関数に実際に変数を与えて処理を行わせるには,以下のように書きます.

(lambda a: a>0 )(1)

このように書けば,変数の1が0より大きいかどうかの判定が行われ,

True

と帰ってきます.

lambda式は式なので当然変数に与えることもできて,

f = lambda a: a>0 

と関数定義のように書くこともできます.この変数fは,map関数のときに普通に定義したfのように使うことができて

f(1)
True

となります.

さて,ここからがlambda式の便利なところです.
これを使えば,上のような0より大きいかどうかを判定する関数を,引数としてmap関数に直接与えることができます.
つまり,1行で書けます.

list(map(lambda a:a>0, [-1,0,1]))

結果は

[False, False, True]

便利ですねえ.


リスト内包表記

任意長の配列

[0, 0, 0, 0, 0, 0, 0]

例えば 上のような長さ7のリストが欲しい時は,

[0]*7

と書けばよいですが,リスト内包表記という”[]の中にfor文を書く”やり方でも作ることができます.

 [0 for i in range(7)]

基本形

 [i for i in range(7)]

出力は

[0, 1, 2, 3, 4, 5, 6]

一番左のiがリストの要素となっています.
この場合では,右側のfor文でiが0から1つずつ増えていくと定義されていることになります.
よって例えば以下のように一番左のiを(i+1)に変えると,

 [(i+1) for i in range(7)]

出力は

[1, 2, 3, 4, 5, 6, 7]

となります.

if(後置)

要素をifで判定してリストに加えるかどうか決めることもできます.
この場合,for文のあとにif文を書きます.

 [i for i in range(7) if i<5]

5より小さいiのみが要素として追加され,

[0, 1, 2, 3, 4]

となります.

if else

if elseを使いたい時は最初にもってくる必要があります.

 [i if i<5 else 4 for i in range(7)]

iが5より小さい時はそのままiが要素となり,iが5以上のときは4が要素となります.

[0, 1, 2, 3, 4, 4, 4]

集計

リストなどで集計関数がそのまま使えます.例えば,

sum( [(i+1) for i in range(7)])

について,全要素の和となるので

28

が得られます.28は完全数ですね.

2重for

forを2回重ねることもできます.

[[i,j] for i in range(2) for j in range(5)]

この場合はiが2通り,jが5通り動くので,長さ10のリストができます.

[[0, 0],
 [0, 1],
 [0, 2],
 [0, 3],
 [0, 4],
 [1, 0],
 [1, 1],
 [1, 2],
 [1, 3],
 [1, 4]]