【NumPy】多次元配列の単一要素を参照するときは括弧を一つで指定した方が速い

はじめに

NumPyで2次元配列の単一の要素を参照するとき、何気なくarray[y][x]のように[][]を二つ使っていませんか？

実は、多次元配列を使うときは、array[y, x]のように[]を一つで書いた方が処理速度が速いです。本記事では、その事について実際に計測して比較しました。

実装

Google Colabで作成した本記事のコードは、こちらにあります。

numpyのインポート

import numpy as np

NumPyのバージョンは1.21.6で実装しています。

結果

Google Colabのセル内をtimeitで計測しました。100ループ(-n 100)を10セット(-r 10)した結果を表記しています。

2次元配列

array[y][x]

%%timeit -r 10 -n 100
array = np.zeros((100, 100))
for y in range(array.shape[0]):
    for x in range(array.shape[1]):
        array[y][x] = 1
# -> 2.73 ms ± 91.9 µs per loop (mean ± std. dev. of 10 runs, 100 loops each)

array[y, x]

%%timeit -r 10 -n 100
array = np.zeros((100, 100))
for y in range(array.shape[0]):
    for x in range(array.shape[1]):
        array[y, x] = 1
# -> 1.74 ms ± 114 µs per loop (mean ± std. dev. of 10 runs, 100 loops each)

	平均
array[y][x]	2.73 ms
array[y, x]	1.74 ms

array[y, x]が一番速い結果になりました。

3次元配列

array[z][y][x]

%%timeit -r 10 -n 100
array = np.zeros((20, 20, 20))
for z in range(array.shape[0]):
    for y in range(array.shape[1]):
        for x in range(array.shape[2]):
            array[z][y][x] = 1
# -> 3.55 ms ± 247 µs per loop (mean ± std. dev. of 10 runs, 100 loops each)

array[z, y][x]

%%timeit -r 10 -n 100
array = np.zeros((20, 20, 20))
for z in range(array.shape[0]):
    for y in range(array.shape[1]):
        for x in range(array.shape[2]):
            array[z, y][x] = 1
# -> 2.78 ms ± 247 µs per loop (mean ± std. dev. of 10 runs, 100 loops each)

array[z][y, x]

%%timeit -r 10 -n 100
array = np.zeros((20, 20, 20))
for z in range(array.shape[0]):
    for y in range(array.shape[1]):
        for x in range(array.shape[2]):
            array[z][y, x] = 1
# -> 2.85 ms ± 261 µs per loop (mean ± std. dev. of 10 runs, 100 loops each)

3次元配列[z, y, x]

%%timeit -r 10 -n 100
array = np.zeros((20, 20, 20))
for z in range(array.shape[0]):
    for y in range(array.shape[1]):
        for x in range(array.shape[2]):
            array[z, y, x] = 1
# -> 1.82 ms ± 275 µs per loop (mean ± std. dev. of 10 runs, 100 loops each)

	平均
array[z][y][x]	3.55 ms
array[z, y][x]	2.78 ms
array[z][y, x]	2.85 ms
array[z, y, x]	1.82 ms

括弧の個数に応じて処理速度が低下し、array[z, y, x]が一番速い結果になりました。

理由

理由としては、aray[y][x]とすると、最初に座標yに属する一次元配列が返され、その返された配列の中でxの座標を指定しています。そのため、一度にaray[y, x]で座標を指定する方が処理速度が速いです。

この事は、NumPy公式に

書かれています。

まとめ

NumPyで多次元配列を使うときは、[]を使うとその都度配列が返されるため処理速度が遅くなることを比較しました。特に3次元配列の結果では、括弧の個数に応じて処理速度が低下していることがわかりました。
NumPyで多次元配列を参照するときは、[]は一つで指定しましょう。

参考資料