目次

Numba Library

競技プログラミングのNumba用スニペットを、作成するたびに追記していく。

基本的には素のPythonとあまり変わらず、せいぜいListをnumpy.ndarrayに置きかえただけのものがほとんどだが、 普通に実装するとNumbaでは使えない機能を踏んでしまうものも一部あり、それを避けた実装としてNumba用にまとまってた方が嬉しい。

クラス変数の代替方法

ちょっと長くなるが、Numba用に移植する際の問題点の1つへの対処法に関する考察を書いておく。

Numbaでは、事前コンパイルをする場合はクラスが使えないのが悩ましい。
おかげでインスタンス変数、つまりは状態を持てないので、常に外部から注入する必要がある。

使う際に何を注入するか意識しなければならないし、管理したい変数が増えると記述もどんどん冗長になる。


class UnionFind:
    def __init(self, n):
        self.table = [-1] * n  # ←インスタンス変数
    
    def unite(self, a, b):
        # ...略

# 使う際にはtableは隠蔽され、中でどう使われているかなんて気にしなくてよい

uft = UnionFind(10)
uft.unite(1, 5)
uft.unite(2, 6)

@njit
def main():
    def unite(table, a, b):
        # 略
    
    def find(table, a, b):
        # 略
    
    # 使う際は、外部でtableを定義し、毎回連れ回す必要が生じる
    
    table = [-1] * 10
    unite(table, 1, 5)
    unite(table, 2, 6)

また、もう1つの問題点として、Numbaは内部関数も含め、nested(入れ子)なリスト等を引数に取れない。Numpyの多次元配列ならOKだが、それでは表現できないものもある。

(上手くいくこともある? 条件調査中)

@njit
def main():
    def something_function(nested_list, a, b):
        # コンパイル時エラー
    
    nested_list = [[0] for _ in range(10)]
    something_function(nested_list, 1, 5)

Numba関数の中で、nestedなリストを作ることはできる。また、同じ関数内にクロージャ関数を作れば、クロージャ関数から関数外のリストを参照することができる。

これを用いて、以下のようにすれば、something_function の中で nested_list が使える。

@njit
def main():

    nested_list = [[0] for _ in range(10)]  # 関数より先にnonlocalな変数を定義

    def something_function(a, b):
        nested_list[a][b] = 1               # 変数を使う
    
    something_function(1, 5)

しかし、それでは関数と状態が1対1で結びついてしまい、クラスにおける「複数のインスタンスを作る」ようなことができない。

あまり綺麗ではないが、無理矢理解決するとしたら、以下のようになるだろうか。

関数外部にはインスタンス別のリストを記録する NESTED_LIST を用意し、init() ではそこに初期化したリストを加える(これが1つのインスタンス変数となる)。
init() は自身のIDを返すので、以降、something_function() など他の関数を呼ぶ際は、そのIDのみを連れ回す。

これなら、管理したい変数が増えても使う側で管理するのはIDのみで済み、極力、内部実装を意識しないで使える。

ちなみに、init()では、入れ子リストの中身が空だと型推定ができず、コンパイルエラーとなってしまう(9~14行目)。
ちょっと奇妙だが、入れる予定の型が分かるような書き方で lst を定義し、それを空にした後コピーするようにすると上手くいく。

@njit
def main():

    NESTED_LIST = []
    
    def init(n):
        _id = len(NESTED_LIST)
        
        # × コンパイルエラー
        # NESTED_LIST.append([[] for _ in range(n)])
        
        lst = [0]
        lst.clear()
        NESTED_LIST.append([lst.copy() for _ in range(n)])
        return _id

    def something_function(_id, a, b):
        nested_list = NESTED_LIST[_id]
        nested_list[a][b] = 1
    
    id1 = something_init(10)
    id2 = something_init(20)
    
    something_function(id1, 1, 5)
    something_function(id2, 2, 6)

あくまでコンパイルが通るというだけで、もっといい書き方があるなら使いたい。

実装例

bit count

2進数表記で'1'の立っている数。


def bit_count(x):
    x = (x & 0x55555555) + ((x >>  1) & 0x55555555)
    x = (x & 0x33333333) + ((x >>  2) & 0x33333333)
    x = (x & 0x0F0F0F0F) + ((x >>  4) & 0x0F0F0F0F)
    x = (x & 0x00FF00FF) + ((x >>  8) & 0x00FF00FF)
    x = (x & 0x0000FFFF) + ((x >> 16) & 0x0000FFFF)
    return x

bit length

2進数表記の桁数(0は0)

    def bit_length(n):
        ret = 0
        while n > 0:
            n >>= 1
            ret += 1
        return ret

なお、$n$ が1以上かを確認するのに while n: としても素のPythonは通るが、Numbaではバージョンにより $n$ をbool値だと推定してコンパイルしてしまう。
そうなると、引数にどんな正整数を渡しても $n=1$ となってしまい、おかしくなる。ちゃんと int 型であることがわかるような書き方をする。

mod累乗

$x^a$ を MOD で割った剰余。pythonなら pow(x, a, MOD) で求められるが、Numbaでは第3引数が未対応。二分累乗法で実装。


def mod_pow(x, a, MOD):
    ret = 1
    cur = x
    while a > 0:
        if a & 1:
            ret = ret * cur % MOD
        cur = cur * cur % MOD
        a >>= 1
    return ret

mod階乗と逆数の事前計算

$0!~N!$ とそのモジュラ逆数を計算。上記の mod_pow を使用。
前提として、$n \lt MOD$ かつ $MOD$ は素数。


    def mod_pow(x, a, MOD):
        ret = 1
        cur = x
        while a > 0:
            if a & 1:
                ret = ret * cur % MOD
            cur = cur * cur % MOD
            a >>= 1
        return ret

    def precompute_factorials(n, MOD):
        factorials = np.ones(n + 1, dtype=np.int64)
        for m in range(2, n + 1):
            factorials[m] = factorials[m - 1] * m % MOD
        inversions = np.ones(n + 1, dtype=np.int64)
        inversions[n] = mod_pow(factorials[n], MOD - 2, MOD)
        for m in range(n, 2, -1):
            inversions[m - 1] = inversions[m] * m % MOD
        return factorials, inversions

Union-Find

union_find_init()table 配列を生成し、返値をインスタンス番号と見なして各関数に与える。
table の根における値は自グループのサイズを表し、結合の際はサイズが小さい方を大きい方の子とする。


    UNIONFIND_TABLE = []

    def unionfind_init(n):
        UNIONFIND_TABLE.append(np.full(n, -1, dtype=np.int64))
        return len(UNIONFIND_TABLE) - 1

    def unionfind_getroot(ins, x):
        table = UNIONFIND_TABLE[ins]
        stack = []
        while table[x] >= 0:
            stack.append(x)
            x = table[x]
        for y in stack:
            table[y] = x
        return x

    def unionfind_unite(ins, x, y):
        table = UNIONFIND_TABLE[ins]
        r1 = unionfind_getroot(ins, x)
        r2 = unionfind_getroot(ins, y)
        if r1 == r2:
            return
        d1 = table[r1]
        d2 = table[r2]
        if d1 <= d2:
            table[r2] = r1
            table[r1] += d2
        else:
            table[r1] = r2
            table[r2] += d1

    def unionfind_find(ins, x, y):
        return unionfind_getroot(ins, x) == unionfind_getroot(ins, y)

    def unionfind_getsize(ins, x):
        table = UNIONFIND_TABLE[ins]
        return -table[unionfind_getroot(ins, x)]

Binary Indexed Tree (Fenwick Tree)

fenwick_init() に要素数を与えて初期化し、返値をインスタンス番号と見なして各関数に与える。

$i$ は $1~N$ の値を取るものとする(0始まりではない)。

lower_boundは、累積和が $x$ 以上になる最小の $i$ を返す。
使わない場合は、FENWICK_LOGN および fenwick_init 内でのそれを求める処理は不要(残しても大した計算量ではないが)。


    FENWICK_TREE = []
    FENWICK_LOGN = []

    def fenwick_init(n):
        log_n = 0
        m = n
        while m:
            log_n += 1
            m >>= 1
        FENWICK_TREE.append(np.zeros(n + 1, dtype=np.int64))
        FENWICK_LOGN.append(log_n)
        return len(FENWICK_TREE) - 1

    def fenwick_add(ins, i, x):
        arr = FENWICK_TREE[ins]
        n = arr.size - 1
        while i <= n:
            arr[i] += x
            i += i & -i

    def fenwick_sum(ins, i):
        arr = FENWICK_TREE[ins]
        result = 0
        while i > 0:
            result += arr[i]
            i ^= i & -i
        return result

    def fenwick_lower_bound(ins, x):
        arr = FENWICK_TREE[ins]
        log_n = FENWICK_LOGN[ins]
        n = arr.size - 1
        sum_ = 0
        pos = 0
        for i in range(log_n, -1, -1):
            k = pos + (1 << i)
            if k < n and sum_ + arr[k] < x:
                sum_ += arr[k]
                pos += 1 << i
        return pos + 1

外部注入できる Fenwick Tree

単位元と演算を外部から注入する版。
ただし、型があまり自由すぎると扱いきれないので、単位元 identity_element の型は np.int64 型固定とし、演算関数 func も「np.int64型の引数を2つとって、1つ返す関数」固定とする。

func は、add, min, xor など operatorモジュールにあるものはそのまま使えるし、自分で定義したものでもよい。

実装

最大流(Dinic法)

辺に容量 $cap_e$ が決められた有向グラフで、頂点 $s$ から $t$ に流せる最大流量を求める。
二部グラフのマッチングにも使える。

頂点番号は $0~N-1$。

基本は、dinic_init で初期化→dinic_add_links でグラフ生成→dinic_maximum_flow で最大流量を計算。

    DINIC_LINKS = []

    def dinic_init(n):
        lst = [[0]]
        lst.clear()
        DINIC_LINKS.append([lst.copy() for _ in range(n)])
        return len(DINIC_LINKS) - 1

    def dinic_add_link(ins, frm, to, cap):
        links = DINIC_LINKS[ins]
        links[frm].append([to, cap, len(links[to])])
        links[to].append([frm, 0, len(links[frm]) - 1])

    def dinic_bfs(ins, n, s):
        links = DINIC_LINKS[ins]
        depth = np.full(n, -1, dtype=np.int64)
        depth[s] = 0
        deq = np.zeros(n + 5, dtype=np.int64)
        dl, dr = 0, 1
        deq[0] = s
        while dl < dr:
            v = deq[dl]
            dl += 1
            for link in links[v]:
                if link[1] > 0 and depth[link[0]] == -1:
                    depth[link[0]] = depth[v] + 1
                    deq[dr] = link[0]
                    dr += 1
        return depth

    def dinic_dfs(ins, depth, progress, s, t):
        links = DINIC_LINKS[ins]
        stack = [(s, 10 ** 18)]
        flow = 0
        while stack:
            v, f = stack.pop()
            if v == t:
                flow = f
                continue
            if flow == 0:
                i = progress[v]
                if i == len(links[v]):
                    continue
                progress[v] += 1
                stack.append((v, f))
                to, cap, rev = links[v][i]
                if cap == 0 or depth[v] >= depth[to]:
                    continue
                stack.append((to, min(f, cap)))
            else:
                i = progress[v] - 1
                link = links[v][i]
                link[1] -= flow
                links[link[0]][link[2]][1] += flow
        return flow

    def dinic_maximum_flow(ins, n, s, t):
        flow = 0
        while True:
            depth = dinic_bfs(ins, n, s)
            if depth[t] == -1:
                return flow
            progress = np.zeros(n, dtype=np.int64)
            path_flow = dinic_dfs(ins, depth, progress, s, t)
            while path_flow != 0:
                flow += path_flow
                path_flow = dinic_dfs(ins, depth, progress, s, t)

最小費用流

辺に容量 $cap_e$ と、1単位を流したときのコスト $cost_e$ が決められた有向グラフで、頂点 $s$ から $t$ に、流量 $Q$ を流した時の最小コストを求める。

頂点番号は $0~N-1$。

基本的な使い方は、mincostflow_init で初期化→mincostflow_add_links でグラフ生成→mincostflow_flow で最小費用を計算。


    from heapq import heappop, heappush

    MINCOSTFLOW_LINKS = []
    INF = 10 ** 10

    def mincostflow_init(n):
        """ n: 頂点数 """
        lst = [[0]]
        lst.clear()
        MINCOSTFLOW_LINKS.append([lst.copy() for _ in range(n)])
        return len(MINCOSTFLOW_LINKS) - 1

    def mincostflow_add_link(ins, frm, to, capacity, cost):
        """ インスタンスID, 辺始点頂点番号, 辺終点頂点番号, 容量, コスト """
        links = MINCOSTFLOW_LINKS[ins]
        links[frm].append([to, capacity, cost, len(links[to])])
        links[to].append([frm, 0, -cost, len(links[frm]) - 1])

    def mincostflow_flow(ins, s, t, quantity):
        """ インスタンスID, フロー始点頂点番号, フロー終点頂点番号, 要求流量 """
        links = MINCOSTFLOW_LINKS[ins]
        n = len(links)
        res = 0
        potentials = np.zeros(n, dtype=np.int64)
        dist = np.full(n, INF, dtype=np.int64)
        prev_v = np.full(n, -1, dtype=np.int64)
        prev_e = np.full(n, -1, dtype=np.int64)

        while quantity:
            dist.fill(INF)
            dist[s] = 0
            que = [(0, s)]

            while que:
                total_cost, v = heappop(que)
                if dist[v] < total_cost:
                    continue
                for i, (u, cap, cost, _) in enumerate(links[v]):
                    new_cost = dist[v] + potentials[v] - potentials[u] + cost
                    if cap > 0 and new_cost < dist[u]:
                        dist[u] = new_cost
                        prev_v[u] = v
                        prev_e[u] = i
                        heappush(que, (new_cost, u))

            # Cannot flow quantity
            if dist[t] == INF:
                return -1

            potentials += dist

            cur_flow = quantity
            v = t
            while v != s:
                cur_flow = min(cur_flow, links[prev_v[v]][prev_e[v]][1])
                v = prev_v[v]
            quantity -= cur_flow
            res += cur_flow * potentials[t]

            v = t
            while v != s:
                link = links[prev_v[v]][prev_e[v]]
                link[1] -= cur_flow
                links[v][link[3]][1] += cur_flow
                v = prev_v[v]

        return res

留意点

同じNumpy配列同士を演算するとエラー

AtCoderで使われて いる いた過去の Numba 0.48.0 では、同じNumPy配列同士を演算すると(?)エラーになる。(※詳細な条件はちゃんと調べてない)

0.53では修正されているのを確認している。

一方を別の名前の変数で定義してやると大丈夫になる。

# エラー(配列 x の各値をそれぞれ a 乗するコード)
@njit
def mod_pow(x, a):
    ret = np.ones_like(x)
    cur = x
    while a > 0:
        if a & 1:
            ret = ret * cur % MOD
        cur = cur * cur % MOD  # ←エラー
        a >>= 1
    return ret

# おっけー
def mod_pow(x, a):
    ret = np.ones_like(x)
    cur = x
    while a > 0:
        if a & 1:
            ret = ret * cur % MOD
        cur_ = cur
        cur = cur * cur_ % MOD
        a >>= 1
    return ret

整数が0か0以外かの判定はちゃんと書く

Numba 0.57.0 で確認。

Pythonでは、整数が0か0以外かの判定に「if a:」「while a:」などと書いても解釈してくれるが、Numbaではbool値として解釈されてしまうことがある。

その場合、コンパイルされた関数中の $a$ は全てbool値となるので、正整数を渡しても強制的に 1 になるなど、おかしくなる。(詳細な条件は不明)


@njit
def main():
    def mod_pow(x, a):
        ret = 1
        cur = x
        while a:        # ← ここの記述からか、a は関数全体を通してbool値として扱われる
            if a & 1:
                ret = ret * cur
            cur = cur * cur
            a >>= 1
        return ret

    print(mod_pow(10, 5))  # 10^1 として渡ってしまい、10 が返る

while a > 0:」など、明示的にint型であることが分かるような書き方をする必要がある。