numpy.ndarrayからインデックスと値の配列を求める

Get values and indexes of an ndarray in Numpy – StackOverflow

「numpy.ndenumerate」を初めて知った.

リターンはイテレータなので,

import numpy as np

A = np.array([[1, 2, 3], [4, 5, 6]])
list(np.ndenumerate(A))

[((0, 0), 1), ((0, 1), 2), ((0, 2), 3), ((1, 0), 4), ((1, 1), 5), ((1, 2), 6)]

求める形にしようと思うと,ループを回す必要がある.

[(*idx, val) for idx, val in np.ndenumerate(A)]

[(0, 0, 1), (0, 1, 2), (0, 2, 3), (1, 0, 4), (1, 1, 5), (1, 2, 6)]
 
 
 
これ,もっと素直に考えれば,[index, value]の形なので,
numpy.argwhereとAのconcatenateを求めれば良いだけじゃないだろうか.
やっている事が単純なので,マルチディメンションでも任意の次元でも関係ない.

np.concatenate([np.argwhere(A), A.reshape(-1, 1)], axis=1)

array([[0, 0, 1],
[0, 1, 2],
[0, 2, 3],
[1, 0, 4],
[1, 1, 5],
[1, 2, 6]], dtype=int64)

numpy.argwhereはnumpy.whereと比べて遅いので,
これがリーズナブルチョイスかどうかは分からないけど.

%timeit [(*idx, val) for idx, val in np.ndenumerate(A)]
%timeit np.concatenate([np.argwhere(A), A.reshape(-1, 1)], axis=1)

10.8 µs ± 341 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
23.1 µs ± 743 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

A = np.ones((4, 4))
%timeit [(*idx, val) for idx, val in np.ndenumerate(A)]
%timeit np.concatenate([np.argwhere(A), A.reshape(-1, 1)], axis=1)

23.7 µs ± 4.3 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
23 µs ± 263 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

A = np.ones((10, 10))
%timeit [(*idx, val) for idx, val in np.ndenumerate(A)]
%timeit np.concatenate([np.argwhere(A), A.reshape(-1, 1)], axis=1)

113 µs ± 7.15 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
29.8 µs ± 3.82 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

ちなみに,

%%timeit
grid = np.mgrid[:A.shape[0], :A.shape[1]]
np.stack([grid[0], grid[1], A]).reshape(3, -1).T

96 µs ± 13.4 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
 
 
 
ちょっとややこしくなるけど,n数が増えた時に速いのは,
色々な場面で多用するnp.arangeを用いてindexerを作るという方法.

rind = np.repeat(np.arange(A.shape[0]), A.shape[1]).reshape(-1, 1)
cind = np.repeat(np.arange(A.shape[1])[None], A.shape[0], axis=0).reshape(-1, 1)
np.concatenate([rind, cind, A.reshape(-1, 1)], axis=1)

array([[0, 0, 1],
[0, 1, 2],
[0, 2, 3],
[1, 0, 4],
[1, 1, 5],
[1, 2, 6]])

%%timeit
rind = np.repeat(np.arange(A.shape[0]), A.shape[1]).reshape(-1, 1)
cind = np.repeat(np.arange(A.shape[1])[None], A.shape[0], axis=0).reshape(-1, 1)
np.concatenate([rind, cind, A.reshape(-1, 1)], axis=1)

25.3 µs ± 3.76 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

A = np.ones((10, 10))
%%timeit
rind = np.repeat(np.arange(A.shape[0]), A.shape[1]).reshape(-1, 1)
cind = np.repeat(np.arange(A.shape[1])[None], A.shape[0], axis=0).reshape(-1, 1)
np.concatenate([rind, cind, A.reshape(-1, 1)], axis=1)

26.5 µs ± 7.32 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

A = np.ones((100, 100))
%timeit [(*idx, val) for idx, val in np.ndenumerate(A)]
%timeit np.concatenate([np.argwhere(A), A.reshape(-1, 1)], axis=1)

14.4 ms ± 1.64 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)
282 µs ± 19.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%%timeit
rind = np.repeat(np.arange(A.shape[0]), A.shape[1]).reshape(-1, 1)
cind = np.repeat(np.arange(A.shape[1])[None], A.shape[0], axis=0).reshape(-1, 1)
np.concatenate([rind, cind, A.reshape(-1, 1)], axis=1)

164 µs ± 10.5 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

広告
カテゴリー: 未分類 パーマリンク

コメントを残す

以下に詳細を記入するか、アイコンをクリックしてログインしてください。

WordPress.com ロゴ

WordPress.com アカウントを使ってコメントしています。 ログアウト / 変更 )

Twitter 画像

Twitter アカウントを使ってコメントしています。 ログアウト / 変更 )

Facebook の写真

Facebook アカウントを使ってコメントしています。 ログアウト / 変更 )

Google+ フォト

Google+ アカウントを使ってコメントしています。 ログアウト / 変更 )

%s と連携中