2023年AtCoder言語アップデートにより、PyPyでもNumPyが使えるようになりました。本記事は、2023年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))
実際の問題を解いてみます。
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]]