【競プロ】PyPy3で使える！NumPy代用ライブラリ

2023年AtCoder言語アップデートにより、PyPyでもNumPyが使えるようになりました。しかしながら、NumPyはオーバーヘッドが大きく、AtCoderで安定的に多次元リスト計算をするには、本記事はまだまだ有用であるようです。

競プロでPythonを使う場合、高速化のためにPyPy3を使うのは常識です。しかし、PyPy3だとNumPyが使えないため、便利なarray計算が使えません。よって、NumPy代用ライブラリを作ってみました。

1. 代用ライブラリ設計の方針

numpy.関数名を、ほぼそのままnp_関数名として再現しており、多重リストやタプルに対して、NumPyと同様な操作を可能にしています。

なお、本物のNumPyでは、1次元の記憶領域をもとに、ビューにより多次元に見せかけている構造になっています。しかし、本ライブラリは、競プロのスニペットとして使うことを目的としているため、多重リスト構造をそのまま利用しています。具体的には以下の設計方針をとっています。

• np_関数名(引数, ...)として利用します。性能・コードの短さと、使いやすさのバランスを取っているため、多重リストまたは多重タプルを直接使うことを前提としています。
• array代用オブジェクトも提供していますが、利用は必須ではありません。
• 関数、引数、引数の型は、代表的なものだけをサポートしています。また多くの関数は1次元または2次元のリストのみをサポートしています。
• 後の方に紹介する関数は、適宜、それ以前に紹介した関数の定義を前提としています。
• _ではじまる関数は、実際のnumpyには無い、内部処理用のものです。
• 例外処理は実装していません。想定外の引数で呼び出された場合、結果は不定です。
• 簡潔にコピーペーストして使えるようにするため、名前付きlambda式を多用したり、短縮表現をしていたりしますので、PEP8ではアラートが出ます。

2. np_関数群

2.1. リスト基本処理

以降のリスト処理の共通的な基本処理です。この節の処理は、$n (\geqq 1)$ 次元リスト全てに適用可能なように作っています。

from functools import reduce
from itertools import zip_longest, accumulate
import operator
import copy
# 基本プロパティ
_is_array_like = lambda a: isinstance(a, tuple) or isinstance(a, list)
np_shape = lambda a: (len(a), *np_shape(a[0])) if _is_array_like(a) else ()
math_prod = lambda a: reduce(operator.mul, a, 1)
np_size = lambda a: math_prod(np_shape(a))
np_ndim = lambda a: len(np_shape(a))
np_strides = lambda a: tuple(accumulate((1, *(np_shape(a)[:0:-1])), operator.mul))[::-1]  # 要素のサイズは1
# 変形
np_flatten = lambda a: np_flatten([a for a in a for a in a]) if np_ndim(a) > 1 else a if np_ndim(a) == 1 else [a]
def _np_shape_complement(shape, size):   # -1を含むshapeをsizeをもとに補完する
    shape0 = shape[0] if shape[0] > 0 else abs(size // math_prod(shape))
    return (shape0,) if len(shape) == 1 else (shape0, *_np_shape_complement(shape[1:], size // shape0))
_np_reshape = lambda a, newshape: [_np_reshape(a[n: n + len(a) // newshape[0]], newshape[1:]) for n in range(0, len(a), len(a) // newshape[0])] if len(newshape) > 1 else a
np_reshape = lambda a, newshape: _np_reshape(np_flatten(a), _np_shape_complement(newshape, np_size(a)))
_np_tile = lambda a, reps2: _np_tile([[copy.deepcopy(a[n + m % reps2[-1][0]]) for m in range(math_prod(reps2[-1]))] for n in range(0, len(a), reps2[-1][0])], reps2[:-1]) if len(reps2) > 0 else a
np_tile = lambda a, reps: _np_tile(np_flatten(a), list(zip_longest(np_shape(a)[::-1], reps[::-1], fillvalue=1))[::-1])[0]
np_broadcast_to = lambda a, shape: np_tile(a, [n // m for m, n in zip_longest(np_shape(a)[::-1], shape[::-1], fillvalue=1)][::-1])
np_broadcast_shapes = lambda *args: tuple(map(max, *[(1,) * (max([len(shape) for shape in args]) - len(shape)) + shape for shape in args]))
np_broadcast_arrays = lambda *args: [np_broadcast_to(a, np_broadcast_shapes(*map(np_shape, args))) for a in args]
# スライス
_np_slice = lambda x, k: x[k] if not _is_array_like(k) else [_np_slice(x, k[1:]) for x in x][k[0]] if len(k) > 1 else x[k[0]] # sはsliceのリスト

2.2. 1次元リスト典型処理

1次元リストに使われる典型処理です。

# 最大値のインデックス
np_argmax = lambda x: max([(x, i) for i, x in enumerate(x)])[-1]
# 最小値のインデックス
np_argmin = lambda x: min([(x, i) for i, x in enumerate(x)])[-1]
# ソート順のインデックス
np_argsort = lambda x: [x[-1] for x in sorted([(x, i) for i, x in enumerate(x)])]

頻出となる累積和はitertools.accumulate()を使うと便利ですが、後述のリスト応用処理としても用意しています。

2.3. 2次元リスト典型処理

2次元リスト（一部例外あり）に使われる典型処理です。一部の関数はリスト基本処理を前提とします。

# 生成
np_zeros = lambda shape: np_reshape([0] * math_prod(shape), shape)  # 任意のn(>= 1)次元に適用可能
np_ones = lambda shape: np_reshape([1] * math_prod(shape), shape)  # 任意のn(>= 1)次元に適用可能
np_identity = lambda n: [[1 if i == j else 0 for j in range(n)] for i in range(n)]   # 2次元のみ
np_arange = lambda *args: list(range(*args))   # 1次元のみ
# 転置・反転・回転　　　# 2次元のみ
np_T = lambda x: [list(x) for x in zip(*x)]
np_flip = lambda x, axis=None: np_flip(np_flip(x, axis=0), axis=1) if axis == None else [x for x in reversed(x)] if axis == 0 else [list(reversed(x)) for x in x]
np_rot90 = lambda x, k=1: np_flip(np_T(x), axis=0) if k % 4 == 1 else np_rot90(np_rot90(x, k=(k + 3) % 4))
# 単項演算(opにはoperatorを入れる)  
# ※ axisは2次元のみ、axis指定なしは任意のn(>= 1)次元に適用可能。
#   ただし、1次元は標準の組み込み関数で代用可能
_np_unary = lambda op, x, axis=None: [op(x) for x in zip(*x)] if axis == 0 else [op(x) for x in x] if axis == 1 or axis == -1 else op(np_flatten(x))
np_sum = lambda x, axis=None: _np_unary(sum, x, axis)
np_max = lambda x, axis=None: _np_unary(max, x, axis)
np_min = lambda x, axis=None: _np_unary(min, x, axis)
np_all = lambda x, axis=None: _np_unary(lambda x: reduce(operator.and_, x) != 0, x, axis)
np_any = lambda x, axis=None: _np_unary(lambda x: reduce(operator.or_, x) != 0, x, axis)
# 二項演算(opにはoperatorを入れる)  ※ 任意のn(>= 1)次元に適用可能
_np_binomial = lambda op, x, y: [_np_binomial(op, x, y) if _is_array_like(x) else op(x, y) for x, y in zip(*np_broadcast_arrays(x, y))]   # ブロードキャスト可
_np_binomial2 = lambda op, x, y: [_np_binomial2(op, x, y) if _is_array_like(x) else op(x, y) for x, y in zip(x, y)]   # ブロードキャスト不可だが高速
np_add = lambda x, y: _np_binomial(operator.add, x, y)
np_sub = lambda x, y: _np_binomial(operator.sub, x, y)
np_mul = lambda x, y: _np_binomial(operator.mul, x, y)
np_floordiv = lambda x, y: _np_binomial(operator.floordiv, x, y)
np_mod = lambda x, y: _np_binomial(operator.mod, x, y)
np_eq = lambda x, y: _np_binomial(operator.eq, x, y)
np_ne = lambda x, y: _np_binomial(operator.ne, x, y)
np_lt = lambda x, y: _np_binomial(operator.lt, x, y)
np_le = lambda x, y: _np_binomial(operator.le, x, y)
np_gt = lambda x, y: _np_binomial(operator.gt, x, y)
np_ge = lambda x, y: _np_binomial(operator.ge, x, y)
np_array_equal = lambda x, y: np_shape(x) == np_shape(y) and np_all(_np_binomial2(operator.eq, x, y))
# インデックス
np_ravel_multi_index = lambda multi_index, shape: np_sum(np_mul(np_T(multi_index), np_strides(np_zeros(shape))), axis=1)
_np_unravel_index = lambda indices, strides:  indices[:-1] if len(strides) == 0 else _np_unravel_index(indices[:-1] + [indices[-1] // strides[0], indices[-1] % strides[0]], strides[1:])
np_unravel_index = lambda indices, shape: np_T([_np_unravel_index([x], np_strides(np_zeros(shape))) for x in indices])
np_nonzero = lambda x: tuple(np_unravel_index([i for i, x in enumerate(np_flatten(x)) if x != 0], np_shape(x)))
# 行列積　　　※ 2次元のみ
from itertools import product
np_matmul = lambda x, y: np_reshape([sum(map(operator.mul, x, y)) for x, y in product(x, np_T(y))], (len(x), -1))

実際の問題を解いてみます。

ABC218_D: Rectangles

N = int(input())
S = [list(input()) for _ in range(N)]
T = [list(input()) for _ in range(N)]

def np_trim_zeros_2d(X):
    x = np_nonzero(np_any(X, axis=0))
    y = np_nonzero(np_any(X, axis=1))
    x1, x2 = min(x[0]), max(x[0])
    y1, y2 = min(y[0]), max(y[0])
    return _np_slice(X, [slice(y1, y2 + 1), slice(x1, x2 + 1)])

S = np_trim_zeros_2d(np_eq(S, '#'))
T = np_trim_zeros_2d(np_eq(T, '#'))
for k in range(4):
    if np_array_equal(S, T):
        print('Yes')
        exit()
    S = np_rot90(S)
print('No')

2.4. リスト応用処理（累積和、座標圧縮）

numpyには無いですが、累積和や座標圧縮といった、競プロで便利に使えるリスト処理を、np_x_関数として用意します。

# ゼータ変換（累積和）、インライン、1〜2次元両用
def np_x_zeta(x, initial=None):
    if isinstance(x[0], list):
        for i in range(len(x)):
            np_x_zeta(x[i], initial=initial)
        if initial is not None:
            x.insert(0, [initial] * len(x[0]))
        for j in range(len(x[0])):
            for i in range(1, len(x)):
                x[i][j] += x[i - 1][j]
    else:
        if initial is not None:
            x.insert(0, initial)
        for i in range(1, len(x)):
            x[i] += x[i - 1]

# メビウス変換（差分）、インライン、1〜2次元両用
def np_x_mebius(x):
    if isinstance(x[0], list):
        for i in range(len(x)):
            np_x_mebius(x[i])
        for j in range(len(x[0])):
            for i in reversed(range(1, len(x))):
                x[i][j] -= x[i - 1][j]
    else:
        for i in reversed(range(1, len(x))):
            x[i] -= x[i - 1]

# 座標圧縮、1〜2次元両用
# 座圧後リスト, 座圧前→後 変換dict, 座圧後→前 変換リスト を返す
import bisect
from collections import defaultdict
def np_x_compress(X, offset=0):
    if isinstance(X[0], list):
        shape = len(X), len(X[0])
        flatten_X = [X for X in X for X in X]
        res, conv, inv = np_x_compress(flatten_X, offset)
        return [res[n * shape[1]: n * shape[1] + shape[1]] for n in range(shape[0])], conv, inv
    else:
        Y = sorted(set(X))
        res, conv, inv = [], defaultdict(int), [None] * (len(Y) + offset)
        for x in X:
            compressed = bisect.bisect_left(Y, x) + offset
            res.append(compressed)
            conv[x] = compressed
            inv[compressed] = x
        return res, conv, inv

3. array代用オブジェクト

前項の各処理を前提として、リストをarrayオブジェクト風に使えるようにするためのラッパーです。よく使う構文のみを提供していますが、ほぼ完全にNumPyの文法で計算することが可能です。また、このarrayオブジェクトはlist型の継承として実装していますので、listに使える関数も直接使えます。

なお、微妙にオーバーヘッドになるため、競プロで使う場合は、array代用オブジェクトは使わずに、np_関数名を直接使うほうがよいでしょう。

class np:
    @classmethod
    def array(cls, a): return ndarray(a)
    @classmethod
    def zeros(cls, shape): return cls.array(np_zeros(shape))
    @classmethod
    def ones(cls, shape): return cls.array(np_ones(shape))
    @classmethod
    def identity(cls, n): return cls.array(np_identity(n))
    @classmethod
    def arange(cls, *args): return cls.array(np_arange(*args))

class ndarray(list):
    def _rc(self, a): return ndarray(a) if _is_array_like(a) else a
    @property
    def shape(self): return np_shape(self)
    @property
    def size(self): return np_size(self)
    @property
    def ndim(self): return np_ndim(self)
    @property
    def T(self): return ndarray(np_T(self))
    def flatten(self): return ndarray(np_flatten(self))
    def reshape(self, *newshape): return ndarray(np_reshape(self, newshape))
    def sum(self, axis=None): return np_sum(self, axis) if axis is None else ndarray(np_sum(self, axis))
    def max(self, axis=None): return np_max(self, axis) if axis is None else ndarray(np_max(self, axis))
    def min(self, axis=None): return np_min(self, axis) if axis is None else ndarray(np_min(self, axis))
    def __add__(self, other): return ndarray(np_add(self, other))
    def __sub__(self, other): return ndarray(np_sub(self, other))
    def __mul__(self, other): return ndarray(np_mul(self, other))
    def __floordiv__(self, other): return ndarray(np_floordiv(self, other))
    def __mod__(self, other): return ndarray(np_mod(self, other))
    def __radd__(self, other): return ndarray(np_add(other, self))
    def __rsub__(self, other): return ndarray(np_sub(other, self))
    def __rmul__(self, other): return ndarray(np_mul(other, self))
    def __rfloordiv__(self, other): return ndarray(np_floordiv(other, self))
    def __rmod__(self, other): return ndarray(np_mod(other, self))
    def __matmul__(self, other): return ndarray(np_matmul(self, other))
    def __getitem__(self, k): return self._rc(_np_slice(list(self), k))

以下が使用例です。NumPyと全く同じ文法で、演算子を組み合わせたり、ブロードキャストを使ったりできています。

x = 5 - np.arange(9).reshape(3, -1)
# [[5, 4, 3], [2, 1, 0], [-1, -2, -3]]
y = x.T @ x
# [[30, 24, 18], [24, 21, 18], [18, 18, 18]]
z = y[0:2, 1:3]
# [[24, 18], [21, 18]]

参考: NumPyの場合

import numpy as np
x = 5 - np.arange(9).reshape(3, -1)
#[[ 5  4  3]
# [ 2  1  0]
# [-1 -2 -3]]
y = x.T @ x
#[[30 24 18]
# [24 21 18]
# [18 18 18]]
z = y[0:2, 1:3]
#[[24 18]
# [21 18]]