競プロ用自作関数置き場

タイトルの通りです。
ここに記載したソースコードは使って頂いても構いませんが、それにより生じた問題や損失についての責任は追いかねます。
おかしなところがあったら教えてもらえると嬉しいです。

全般

再帰関数をやるときのおまじない

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.par[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]]