numpyの二次配列があって、次のように行と列をindexingすることを考えます。
(下図では6x6の二次配列から[[9, 12], [21, 24]]を取り出そうとしています)
このとき、計算方法によっては速度にかなりの差がでました!
# 巨大な二次配列からランダムな行、列の要素を抽出する
import numpy as np
N = 10000
X = np.arange(N ** 2).reshape(N, N)
M = 100
a = np.random.choice(N, M)
b = np.random.choice(N, M)
%timeit Y1 = X[a][:, b]
# 実行例) 1.09 ms ± 243 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit Y2 = X[a[:, np.newaxis], b]
# 実行例) 66.8 µs ± 1.56 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
X[a][:, b]よりもX[a[:, np.newaxis], b]の計算のほうが圧倒的に高速です。
以下は蛇足ですが、そもそも上の図のように二次配列に対してindexingするのに手こずったので、ここに至るまでの紆余曲折を残しておきます。
Sliceの場合
この場合は簡単で
import numpy as np
N = 10000
X = np.arange(N ** 2).reshape(N, N)
a = np.s_[0:100] # 行に関するslice
b = np.s_[100:300] # 列に関するslice
# スライスした結果
Y = X[a, b] # Y.shape = (100,200)
と表現出来ます。ただ、もっと柔軟に行・列を取り出そうとするとスライスだと限界があります。
Indexingの場合 (失敗例)
上と同じようにやろうとすると実はうまくいきません。
import numpy as np
N = 10000
X = np.arange(N ** 2).reshape(N, N)
a = np.arange(0,100) # 行に関するindexing
b = np.arange(100, 300) # 列に関するindexing
# indexingの結果
Y = X[a, b] # エラーが起きる
二次元配列に対して上記のようにindexingする場合、aとbは同じ長さの配列でなければなりません。
また、同じ長さでも次のような配列が返ってきます。
import numpy as np
N = 10000
X = np.arange(N ** 2).reshape(N, N)
a = np.arange(3) # 行に関するindexing
b = np.arange(3) # 列に関するindexing
# indexingの結果
Y = X[a, b] # 結果: [0, 10001, 20002]
ちょうどY[i] = X[a[i], b[i]]といった挙動です。
Indexing (成功例)
上の方法が駄目ということで以下の通り書き換えると想定通りにindexingできました。
import numpy as np
N = 10000
X = np.arange(N ** 2).reshape(N, N)
a = np.arange(0,100) # 行に関するindexing
b = np.arange(100, 300) # 列に関するindexing
# indexingの結果
Y = X[a][:, b]
ただ、numpyのドキュメントを読んでみるとこんな記述が。
So note that x[0,2] = x[0][2] though the second case is more inefficient as a new temporary array is created after the first index that is subsequently indexed by 2.
つまりX[a]とすることで、一時的な配列が生成されて効率が悪いようです。
実際、冒頭のコードでも見せたように二次配列が大きくなると計算時間に如実に影響を与えました。
再掲ですが最終的なindexing方法は以下のとおりです。
import numpy as np
N = 10000
X = np.arange(N ** 2).reshape(N, N)
a = np.arange(0,100) # 行に関するindexing
b = np.arange(100, 300) # 列に関するindexing
# indexingの結果
Y = X[a[:, np.newaxis], b]
以上、蛇足でした。
試行錯誤での結果なので、もっと効率的な方法があれば是非教えていただきたいです。