タイトルの通りです。
ここに記載したソースコードは使って頂いても構いませんが、それにより生じた問題や損失についての責任は追いかねます。
おかしなところがあったら教えてもらえると嬉しいです。
全般
再帰関数をやるときのおまじない
import sys
sys.setrecursionlimit(10**8)
累積和
reverse = Trueにすると、後ろからの累積和
def cumsum(l, reverse = False):
if not reverse:
cum = 0
cumsum_l = []
for ll in l:
cum += ll
cumsum_l.append(cum)
return cumsum_l
else:
return list(reversed(cumsum(reversed(l))))
整数関連
進数を変える関数
bef進数で表された数m(str型)をaft進数にする。
使ってもいい文字は0~9だけです。16進数のA~Fみたいなのが来ると困ります。
base_change.py
def base_change(m,bef,aft):
if m == "0":
return "0"
p = sum(int(m[i]) * bef ** (len(m) - i - 1) for i in range(len(m)))
q = ""
while(p != 0):
q = str(p % aft) + q
p //= aft
return q
十進数をN進数に変換する
負の進数(-2進数など)にも対応
def negative_div(a,b):
if b > 0:
return a//b, a%b
elif a%b == 0:
return a//b, 0
else:
return a//b+1, a%b-b
def change_base(n, base):
after = []
while n != 0:
a,b = negative_div(n,base)
after.append(b)
n = a
return after
使い方
print(change_base(32, 2)) -> [0, 0, 0, 0, 0, 1]
print(change_base(314159, 265)) -> [134, 125, 4]
print(change_base(12345, -7)) -> [4, 1, 0, 6, 6]
1からnまでの整数のxor
https://atcoder.jp/contests/abc121/tasks/abc121_d
この問題が解けます
prime_list.py
def xor_sum(n):
ans = 0
keta = len(bin(n))-2
n += 1
for i in range(keta):
block = 1<<(i+1)
ans += (((n//block) * (block // 2) + max(0,n % block - block // 2)) & 1) << i
return ans
素数列挙
prime_list.py
def prime_list(n):
pl = []
is_prime = [True] * (n+1)
is_prime[0] = False
is_prime[1] = False
for i in range(2,n+1):
if is_prime[i]:
j = 2
while(i*j <= n):
is_prime[i*j] = False
j += 1
for i,b in enumerate(is_prime):
if b:
pl.append(i)
return pl
使い方↓
print(prime_list(100)) # [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97]
素因数分解(その1)
dict型で返します。
計算量は $ O(\sqrt{n}) $ 。
prime_factorization.py
def prime_factorization(n):
a = 2
ans = {}
m = n
while(a**2 <= n):
if m%a == 0:
if a in ans:
ans[a] += 1
else:
ans[a] = 1
m //= a
else:
a += 1
if m != 1:
ans[m] = 1
return ans
インプットとアウトプット↓
print(prime_factorization(720)) -> {2: 4, 3: 2, 5: 1}
print(prime_factorization(169)) -> {13: 2}
print(prime_factorization(7)) -> {7: 1}
素因数分解(大量処理するときに速いバージョン。約数列挙つき)
はじめに2以上max_n(素因数分解対象の整数の最大値)の素因数のうち最小のものを計算してminfactorに格納。nを素因数分解するときにはminfactor[n]でnを割り、その商(n')をさらにminfactor[n']で割り・・・というのを1になるまで繰り返す。割った値は素因数なのでを記録しておき素因数分解の結果とする。
計算量はminfactorの算出で $ O(NloglogN) $ 、素因数分解の実行で $ O(logN) $ 。
参考文献：
https://qiita.com/drken/items/3beb679e54266f20ab63
prime_factorization2.py
class FastQueryPrimeFactorization:
def __init__(self,max_n) -> None:
self.max_n = max_n
if not isinstance(self.max_n,int) or self.max_n <= 1:
raise ValueError("max_n must be an integer greater than 1")
self.minfactor = [0]*(self.max_n+1)
p = 2
while p <= self.max_n:
i = 1
while(p*i <= self.max_n):
if self.minfactor[p*i] == 0:
self.minfactor[p*i] = p
i += 1
p += 1
def prime_factrization(self,n) -> dict:
if not isinstance(n,int) or n <= 1 or n > self.max_n:
raise ValueError("n must be an integer between 2 and max_n({})".format(self.max_n))
prime_factor = {}
while n != 1:
if self.minfactor[n] not in prime_factor:
prime_factor[self.minfactor[n]] = 1
else:
prime_factor[self.minfactor[n]] += 1
n //= self.minfactor[n]
return prime_factor
def divisor(self, n):
from itertools import product
if not isinstance(n,int) or n > self.max_n:
raise ValueError("n must be an integer between 1 and max_n({})".format(self.max_n))
if n == 1:
return [1]
devisor = []
primes = self.prime_factrization(n)
primes_list = [p for p in primes]
expons = product(*(range(primes[p]+1) for p in primes))
for expon in expons:
d = 1
for i, p in enumerate(primes_list):
d *= p ** expon[i]
devisor.append(d)
devisor.sort()
return devisor
使い方↓
f = FastQueryPrimeFactorization(100000)
print(f.prime_factrization(31415)) # => {5: 1, 61: 1, 103: 1}
print(f.prime_factrization(92653)) # => {11: 1, 8423: 1}
print(f.divisor(120)) # => [1, 2, 3, 4, 5, 6, 8, 10, 12, 15, 20, 24, 30, 40, 60, 120]
print(f.divisor(100)) # => [1, 2, 4, 5, 10, 20, 25, 50, 100]
二項係数のmod
class ModConbination:
def __init__(self, n_max, mod) -> None:
self.mod = mod
self.frac = [1]
self.frac_inv = [pow(1,-1,self.mod)]
for i in range(1,n_max+1):
self.frac.append(self.frac[-1] * i % self.mod)
self.frac_inv.append(self.frac_inv[-1] * pow(i,-1,self.mod) % self.mod)
def combination(self, n, k):
if k > n:
return 0
if k < 0 or n < 0:
return 0
else:
return self.frac[n] * self.frac_inv[k] * self.frac_inv[n-k] % self.mod
def h_combination(self, n, k):
return self.combination(n+k-1, k)
行列
行列の掛け算をする関数(mat_dot)
行列の掛け算はnumpyでできるのですが、要素の値が大きくなるとオーバーフローするのでlistで管理
inputもoutputもlist形式です
(@shiracamus様より高速かつシンプルなコードをいただきましたので、そちらを掲載しています。ありがとうございます！)
mat_dot.py
def mat_dot(mat1, mat2):
if len(set(map(len, mat1))) != 1 or len(set(map(len, mat2))) != 1:
raise ValueError("wrong matrix")
if len(mat1[0]) != len(mat2):
raise ValueError("matrix size dose not match")
return [[sum(v1 * v2 for v1, v2 in zip(r1, r2)) for r2 in zip(*mat2)]
for r1 in mat1]
行列の掛け算をして、各要素をmodで割った剰余にする関数(mat_dot_mod)
inputもoutputもlist形式です
mat_dot_mod.py
def mat_dot_mod(mat1, mat2, mod):
if len(set(map(len, mat1))) != 1 or len(set(map(len, mat2))) != 1:
raise ValueError("wrong matrix")
if len(mat1[0]) != len(mat2):
raise ValueError("matrix size dose not match")
return [[sum(v1 * v2 for v1, v2 in zip(r1, r2)) % mod for r2 in zip(*mat2)]
for r1 in mat1]
行列の累乗の剰余を出す関数
組み込み関数pow(a,p,mod)の行列版
上のmat_dot_modと一緒にコピペしないとエラーします
mat(第1入力変数)とoutputはlist形式です
mat_pow_mod.py
def mat_pow_mod(mat,p,mod):
c = mat
len_mat = len(mat)
ans = [[1 if i == j else 0 for j in range(len(mat))]for i in range(len(mat))]
for i in range(len(bin(p))-2):
if 1 << i & p:
ans = mat_dot_mod(ans,c,mod)
c = mat_dot_mod(c,c,mod)
return ans
グラフ
UnionFind木
蟻本準拠ですが少し工夫をプラスして実装しました。
【使い方】
find(x)でxの根を求める。
unite(x,y)でxの属する集合とyの属する集合を結合する。
xやyがなければ、それぞれ自身を根とした集合を新設する。
same(x,y)でxとyが同じ集合か確認する。
xとyのいずれかがまだuniteに登場していないものだった場合はFalseを返します。
connected_components()で連結成分のリストを出す。
dict型で親の管理をしているので数値だけでなくタプル(正確にはハッシュ可能なもの)もxやyに入れられるようにしています。
UnionFindTree.py
class UnionFindTree():
def __init__(self):
self.par = {}
self.rank = {}
self._size = {}
def find(self,x): #根の計算
if self.par[x]==x:
return x
else:
self.par[x] = self.find(self.par[x])
return self.par[x]
def unite(self,x,y): #集合の結合
if x in self.par:
x = self.find(x)
else:
self.par[x] = x
self.rank[x] = 0
self._size[x] = 1
if y in self.par:
y = self.find(y)
else:
self.par[y] = y
self.rank[y] = 0
self._size[y] = 1
if x == y:
return
if self.rank[x]<self.rank[y]:
self.par[x] = y
self._size[y] += self._size[x]
elif self.rank[x]==self.rank[y]:
self.par[y] = x
self.rank[x] += 1
self._size[x]+= self._size[y]
else:
self.par[y] = x
self._size[x]+= self._size[y]
def same(self,x,y): #同じ集合か判定
if x not in self.par or y not in self.par:
return False
else:
return self.find(x)==self.find(y)
def size(self, n): #nが属する集合のサイズを返す
return self._size[self.find(n)]
def connected_components(self): #連結成分
cc_list = {}
for p in self.par:
par_p = self.find(p)
if par_p in cc_list:
cc_list[par_p].add(p)
else:
cc_list.update({par_p:set([p])})
return list(cc_list.values())
使用例
uf = UnionFindTree()
uf.unite(1,2)
uf.unite(2,3)
uf.unite(4,3)
uf.unite(5,6)
print(uf.connected_components()) #[{1, 2, 3, 4}, {5, 6}]
print(uf.same(1,6)) #False
重みつきUnion Find
ダイクストラ法
使い方
初期化：directed = Trueと指定すると有向グラフ扱いできる。(デフォルトはFalse:無向グラフ)
set_path(start,terminal,cost)：始点と終点とコストを設定する。無向グラフではstartとterminalはどっちが先でも同じ。
excecute(n)：nを始点として各点への最低コストを算出する。
dijkstra.py
import heapq
class dijkstra():
def __init__(self,directed=False):
self.directed = directed
self.g = {}
self.v_list = set()
def set_path(self,s,t,c):
self.v_list.add(s)
self.v_list.add(t)
if s in self.g:
self.g[s].append((c,t))
else:
self.g[s] = [(c,t)]
if not self.directed:
if t in self.g:
self.g[t].append((c,s))
else:
self.g[t] = [(c,s)]
def execute(self,start):
self.d = {}
for v in self.v_list:
self.d[v] = float("inf")
self.d[start] = 0
que = []
heapq.heappush(que,(0,start))
while que:
q = heapq.heappop(que)
if q[0] > self.d[q[1]]:
continue
if q[1] not in self.g:
continue
for e in self.g[q[1]]:
cost_n, point_n= e[0],e[1]
if self.d[point_n] > q[0] + cost_n:
self.d[point_n] = q[0] + cost_n
heapq.heappush(que,(self.d[point_n],point_n))
return self.d
使用例
dij = dijkstra()
dij.set_path(1,2,2)
dij.set_path(1,3,3)
dij.set_path(2,5,2)
dij.set_path(3,4,1)
dij.set_path(3,5,4)
dij.set_path(4,7,5)
dij.set_path(5,6,1)
dij.set_path(5,7,6)
dij.set_path(6,7,3)
print(dij.execute(1)) # {1: 0, 2: 2, 3: 3, 4: 4, 5: 4, 6: 5, 7: 8}
連結成分分解
set_path(start,terminal)：始点と終点を設定する。
excecute()：連結成分の数とラベルを返す。
scipyを使っているのでPyPyで出すとおこられる。
connected_components.py
from scipy.sparse.csgraph import connected_components
from scipy.sparse import csr_matrix
class ConnectedComponents:
def __init__(self, n):
"""
n : number of nodes
"""
self.n = n
self.v = []
self.i = []
self.j = []
def set_path(self, s, t):
"""
s : start point
t : terminal point
弱連結成分分解を利用する場合は順不同。
"""
self.v.append(1)
self.i.append(s)
self.j.append(t)
def execute(self, connection = "weak"):
"""
connection : ['weak'|'strong'] the type of connection to use
"""
csr = csr_matrix((self.v, (self.i, self.j)),shape=(self.n + 1, self.n + 1))
n_groups, labels = connected_components(csr, connection = connection, return_labels = True)
groups = {i:[] for i in range(n_groups-1)}
for i in range(self.n):
groups[labels[i+1] - 1].append(i+1)
return groups
使い方
con = ConnectedComponents(5)
con.set_path(1,2)
con.set_path(2,5)
con.set_path(3,4)
print(con.execute()) # -> {0: [1, 2, 5], 1: [3, 4]})
↓自作バージョン。再帰関数ではなくlistをstackとした非再帰DFSで実装したのでpypyでもそれなりに速い。
非再帰DFSのアイデアはこちらから
scc.py
from collections import defaultdict
class SCC:
def __init__(self, n):
self.n = n
self.path = defaultdict(set)
self.path_inv = defaultdict(set)
def set_path(self, s, t):
self.path[s-1].add(t-1)
self.path_inv[t-1].add(s-1)
def execute(self):
sccs = dict()
seen1 = [False] * self.n
t = []
for i in range(self.n):
if seen1[i]:
continue
init_u = i
stack = [(init_u, True)]
while stack:
u, preorder = stack.pop()
if preorder:
if seen1[u]:
continue
seen1[u] = True
stack.append((u, False))
for v1 in self.path[u]:
if not seen1[v1]:
stack.append((v1, True))
else:
t.append(u)
t.reverse()
seen2 = [False] * self.n
scc_id = 0
for j in range(self.n):
if seen2[t[j]]:
continue
scc_tmp = [t[j]+1]
stack = [t[j]]
seen2[t[j]] = True
while stack:
u = stack.pop()
for v2 in self.path_inv[u]:
if not seen2[v2]:
stack.append(v2)
scc_tmp.append(v2+1)
seen2[v2] = True
sccs[scc_id] = scc_tmp
scc_id += 1
return sccs
使用例
scc1 = SCC(7)
scc1.set_path(1,2)
scc1.set_path(2,3)
scc1.set_path(3,1)
scc1.set_path(3,4)
scc1.set_path(4,5)
scc1.set_path(5,4)
scc1.set_path(5,6)
scc1.set_path(5,7)
scc1.execute() # {0: {1, 2, 3}, 1: {4, 5}, 2: {7}, 3: {6}}
平衡二分木的なもの
👉 https://github.com/tatyam-prime/SortedSet/blob/main/SortedSet.py
ソートのルール変更
abc308cなどに対応
要素数2の有理数のlist([分子,分母])をソートする。(有理数のソートならFractionsを使った方がいいかも)
from functools import cmp_to_key
def cmp(a, b):
if a[0] * b[1] == b[0] * a[1]:
return 0
return -1 if a[0] * b[1] < b[0] * a[1] else 1
l = [[999999999, 1000000000],[333333333, 999999999],[1000000000, 999999997],[999999998, 1000000000],]
print(sorted(l, key=cmp_to_key(cmp), reverse = True)) #[[1000000000, 999999997], [999999999, 1000000000], [999999998, 1000000000], [333333333, 999999999]]
商の剰余
$N$が$M$で割り切れる時、
\frac{N}{M}\ \mathrm{mod}\ P = \frac{N\ \mathrm{mod}\ MP}{M}
が成り立ちます。
これを使うと、ABC357Dなどは簡単に書けます。
n = int(input())
mod = 998244353
p = len(str(n))
print(n * ((pow(10, p * n, mod * (10**p - 1))) // (10**p - 1)) % mod)
簡単な説明
$N$が$M$で割り切れることから、$N'$を整数として$\frac{N}{M}$=$N'$と表すことができます。
ここで、$X$および$Y$をそれぞれ$N'$を$P$で割った商と余りとした時、
N' = PX + Y
が成り立ちます。求めたいものは$Y$です。
両辺に$M$を掛けると
N = N'M = PMX + YM
この式から、$N$を$PM$で割った余りは$YM$です。
したがって、$Y$を求めるためには、この$YM$を$M$で割ればいいです。
すなわち、$\frac{N}{M}$を$P$で割った余りは、$N$を$PM$で割った余りを$M$で割れば得られます。