競プロメモ Python編(便利な機能や手抜き方法)

小数点計算にはDecimalを使う

競プロで浮動小数点の罠系の問題が出る場合、Python以外にも解ける解法があるはず。という点は留意しておこう。ただし、何も考えずに高い精度が出せるPythonのDecimalは便利だ。

尚、Pythonといえど、普通に使うと浮動小数点の罠は生じる。例えば、0.35を小数第一位まで四捨五入したいが以下のように0.3となってしまうDecimalを使おう。quantizeは少し使う方が特殊で、以下のように"0.001"などの表示になるように四捨五入、切り下げ、切り上げをしてくれる。

# 0.35: 0.3 になる！四捨五入してくれない
print('0.35:', round(0.35, 1))

from decimal import Decimal
import decimal

# 標準的な使い方。Decimal同士の計算
x = Decimal(1) / Decimal(3) # 0.3333333333333333333333333333
print(type(x)) # <class 'decimal.Decimal'>

# quantize
f = "0.35"
"""
> 0.35
> 0.4
> 0.3
"""
print(Decimal(f).quantize(Decimal('0.01'), rounding=decimal.ROUND_HALF_UP))
print(Decimal(f).quantize(Decimal('0.1'), rounding=decimal.ROUND_CEILING))
print(Decimal(f).quantize(Decimal('0.1'), rounding=decimal.ROUND_FLOOR))

sortだけじゃない。max,minへのkey指定

sortはkeyが指定できます。max, minにもkeyを指定できます。

dat = [ ("c", 10), ("a", 30), ("b", 20), ]
# datの[1]を各種キーにしたい
# まずはsort
print(sorted(dat, key=lambda x: x[1])) # [('c', 10), ('b', 20), ('a', 30)]
print(sorted(dat, key=lambda x: x[1], reverse=True)) # [('a', 30), ('b', 20), ('c', 10)]
#  max, minでも実は指定できる
print(max(dat)) # ('c', 10)
print(max(dat, key=lambda x: x[1])) # ('a', 30)
print(min(dat, key=lambda x: x[1])) # ('c', 10)

ソートの複数キー指定

sortは複数のキーで並べ替えられます。keyをタプルやリストを返すlambda式で指定します。
尚、基本は昇順なので、[1]を昇順->[0]を降順などしたい場合、lambda内でよしなに処理しましょう。

dat = [ (1, 10), (2, 30), (3, 20), (2, 20), (2, 10)]
# [1]の昇順にした上で、[1]が一緒なら[0]で並び買える
dat.sort(key=lambda x: [x[1], x[0]] )
print(dat) # [(1, 10), (2, 10), (2, 20), (3, 20), (2, 30)]
# [1]の昇順にした上で、[1]が一緒なら[0]で【降順に】並び買える
dat.sort(key=lambda x: [x[1], -x[0]] )
print(dat) # [(2, 10), (1, 10), (3, 20), (2, 20), (2, 30)]

modIntのpow

$a^b mod m$ はpowの第三引数で与えることでできます。これは、ビルトインのpowに限られmath.powは第三引数がないので注意。
尚、bが大きくてもO(b)じゃなくて、ちゃんと最適化した計算量でやってくれます。

# 2^3 の mod 5
print(pow(2, 3, 5)) # 3

anyとall は mapと組み合わせよう

any, allは例えば、「偶数が配列に含まれるか？」「全部偶数か？」というTrue, Falseを判定するのに便利。
動作としては、要素が1つでも/すべて Trueの時にTrueを返す。

引数の配列と書いたが、これはiterで渡すので、mapを与えることができる。

dat = [2,4,5,6]
print(dat)
print(any(map(lambda x: x%2 == 0, dat))) # True
print(all(map(lambda x: x%2 == 0, dat))) # False
dat = [2,4,6]
print(dat)
print(any(map(lambda x: x%2 == 0, dat))) # True
print(all(map(lambda x: x%2 == 0, dat))) # True

divmodで商と余りを同時に求める

divmodを使うと商と余りを同時に求められる

# q, r = 10//3, 10%3 と等価
q, r = divmod(10, 3)
print(q,r) # 3 あまり 1

enumerateのスタートのインデックスを指定(1 origin)

enumerateをつかうと　配列に対してindexをつけられるが、0 originである。

print(list(enumerate(["a", "b", "c", "b"])))
# [(0, 'a'), (1, 'b'), (2, 'c'), (3, 'b')]

時折、1 originの出題があるが、この際、startを指定してやると、1 originにできる。

dat = ["a", "b", "c", "b"] # の中から、bのindexを1 originで求める
buf = list(enumerate(dat, start=1))
print(buf) # [(1, 'a'), (2, 'b'), (3, 'c'), (4, 'b')]
res = list(filter(lambda x: x[1] == "b", buf)) # [(2, 'b'), (4, 'b')]
res = list(map(lambda x: x[0], res))
print(res) # [2, 4]

asciiコードを調べる<->asciiコードから文字にする

あんまり良い方法はないので愚直に。

print(ord("a")) # 97
print(hex(ord("a"))) # 0x61
print(hex(ord("a"))[2:])  # 61
print(chr(97)) # a
# 0x61 -> 61 に sliceしてから10進数に変換
print(chr(int("0x61"[2:], 16))) # a

evalで文字列の数式を評価する

evalは"1+2"のような文字列を評価して3を返す。尚、Python上での変数を扱うこともできる。

print(eval("1+2*2")) # 5
a = 2
b = 3
print(eval("a*b")) # 6

n進数の処理

Intは第二変数にbaseを指定できる。

print(int("ff", 16)) # 16進数 ff > 255
print(int("a1", 11)) # 11進数 a1 > 111
print(int("11", 7)) # 7進数 11 > 8

listのdeepcopy

Python使いが１度ははまったことのあるはずの、リストは参照でのコピー。deepcopyを使うと、値をコピーできる。もちろん、$O(N)$なので注意ですが、計算の少ない時に自分で複製するより確実で早い。

l = [1,2,3]
lnew = l
lnew[0] = 0
print(l) # [0,2,3]
###
from copy import deepcopy
l = [1,2,3]
lnew = deepcopy(l)
lnew[0] = 0
print(l) # [1,2,3]

同じ文字の繰り返し数を調べるのに便利な itertools.groupby

[list(g) for k, g in groupby('AAAABBBCCD')] --> AAAA BBB CC D

という例があるが、Nimなどの場合、少しwrapするとすごく使いやすくなる。

"""
>>> countstrs("1110000111")->    [('1', 3), ('0', 4), ('1', 3)]
>>> countstrs("") ->    []
>>> countstrs("aaa") ->    [('a', 3)]
>>> countstrs_withIndex("1110000111")-> [('1', 3, 0), ('0', 4, 3), ('1', 3, 7)]
"""
import itertools

def countstrs(s):
    return [(k, len(list(g))) for k, g in itertools.groupby(s)]
def countstrs_withIndex(s):
    d = countstrs(s)
    r = []
    ind = 0
    for i in range(len(d)):
        r.append((d[i][0], d[i][1], ind))
        ind += d[i][1]
    return r

文字列処理で必須 collection.Counter

文字列など、要素を高速にカウントする。
most_commonを調べると最大の文字数の文字を調べられる。

注意: 同順、つまり、一番多い文字が複数ある場合でも、most_common(1)は1つしか要素を返さない

from collections import Counter
s = "aaabbbcddfaaa"
c = Counter(s)
print(c) # Counter({'a': 6, 'b': 3, 'd': 2, 'c': 1, 'f': 1})
print(c.most_common()) # Counter({'a': 6, 'b': 3, 'd': 2, 'c': 1, 'f': 1})

# counterは引き算ができる
s1 = "21111000"
s2 = "311110"
c1 = Counter(s1) # Counter({'1': 4, '0': 3, '2': 1})
c2 = Counter(s2) # Counter({'1': 4, '3': 1, '0': 1})
print(c1.subtract(c2))
print(c1) # Counter({'0': 2, '2': 1, '1': 0, '3': -1})

# subtractはCounterそのものを書き換える。"-"演算をすると少し挙動が違う
c1 = Counter(s1) # Counter({'1': 4, '0': 3, '2': 1})
c2 = Counter(s2) # Counter({'1': 4, '3': 1, '0': 1})
print(c1-c2) # Counter({'0': 2, '2': 1}) -1の要素が存在しなくなっている

collection.defaultdictの初期値を指定する

defaultdictは便利だが、初期値を変える際はlambda式(引数なし)が必要になる。

from collections import defaultdict
d = defaultdict(int)
d["hoge"] += 10
print(d)
d = defaultdict(lambda : int(10))
d["hoge"] += 10
print(d)

TIPS: collection.deque のrotate

左右にシフトすることができる。

from collections import deque
q = deque([1,2,3,4,5])
print(q)
q.rotate(1) # 右にシフト
print(q)
q.rotate(-2) # 左にシフト
print(q)

ZIP

あんまり競プロで使う場面に当たったことがないのですが、複数の配列の同じ個数目を結合します。unzipはzip(*)

a = [1, 2, 3]
b = ["a", "b", "c"]
l = list(zip(a,b))
print(l) # [(1, 'a'), (2, 'b'), (3, 'c')]
l = list(map(lambda x: list(x), l))
print(l) # [[1, 'a'], [2, 'b'], [3, 'c']]

# 復元(ばらす)
aa, bb = zip(*l)
print(aa,bb) # (1, 2, 3) ('a', 'b', 'c')

# zipは2つである必要はない
a = [1, 2, 3]
b = [4, 5, 6]
c = [7, 8, 9]
print(list(zip(a, b, c))) # [(1, 4, 7), (2, 5, 8), (3, 6, 9)]

# 逆も然り。こういうのをばらせる
dat = [(1,2,3), (4,5,6), (7,8,9)]
aa, bb, cc = zip(*dat)
print(aa,bb,cc) # (1, 4, 7) (2, 5, 8) (3, 6, 9)

[汎用]ソートされている配列の特定範囲の個数を計算

別にこれはPython固有のテクニックではないですね。ソートされている配列において特定範囲の個数を求めたい場合、二分探索が使える。

dat = [1,2,3,4,4,4,4,4,5,6,7,8]
# [3,6]の個数
cnt = bisect.bisect_right(dat, 6) - bisect.bisect_left(dat,3)
print(cnt)
# bisect_leftを両方に使う場合は[3, 6+1)となる
cnt = bisect.bisect_left(dat, 6+1) - bisect.bisect_left(dat,3)
print(cnt)

itertools

import itertools
dat = [1,3,7,3]

# groupby ものすごく使える
#[k for k, g in groupby('AAAABBBCCDAABBB')] --> A B C D A B
#[list(g) for k, g in groupby('AAAABBBCCD')] --> AAAA BBB CC D

# accumulate: 順に演算。max等も取れるので、ケース次第。
# [1, 3, 21, 63]
print(list(itertools.accumulate(dat, lambda x,y: x*y)))
# [1, 3, 7, 7]
print(list(itertools.accumulate(dat, max)))

# chain: リストをイテレータ的に結合する。メモリを食わない。
# [1, 2, 3, 4, 5]
print(list(itertools.chain([1,2], [3,4], [5])))
# 累積和取るのに使えるかも？ [0, 1, 3, 6, 10]
print(list(itertools.accumulate(itertools.chain([0], [1, 2, 3, 4]))))

# 有用度: 問題次第
# dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1 条件を満たしたところからの配列を返す。

# repeat 同じ値を繰り返す。
# ケースにより使えそう。例えば、i^2のリストを生成できる。
#[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
print(list(map(pow, range(10), itertools.repeat(2))))

# みんな大好きPとC, 便利なんだけど、実際あまり必要になったことがない
# permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
# permutations(range(3)) --> 012 021 102 120 201 210
# combinations('ABCD', 2) --> AB AC AD BC BD CD
# combinations(range(4), 3) --> 012 013 023 123
# combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC

# product 使いどころが難しい
# [(1, 2), (1, 3), (1, 4), (2, 2), (2, 3), (2, 4), (3, 2), (3, 3), (3, 4)]
print(list(itertools.product([1,2,3], [2,3,4])))