Numpy配列の非ゼロ要素のインデックスを取得する

とある大き目のマトリックスを対象にゼロでない要素のインデックス(行番号と列番号)を取得する必要があったので、
スマートな方法を探しました。
全部の行と列を2重のfor文で探索すればできるのですがあまりにも無駄が多い気がしましたので。

結果的に、nonzeroという、ずばりそのままの名前の関数があったのでそれを使っています。
ドキュメント: numpy.nonzero

まず、1次元の配列に対してやって見ましょう。

import numpy as np
ary1 = np.array([0, 0, 5, 3, 0, -1, 0])

# 非ゼロ要素のインデックスを取得する
print(np.nonzero(ary1))
# (array([2, 3, 5]),)

# 別の書き方
print(ary1.nonzero())
# (array([2, 3, 5]),)

# 非ゼロ要素の値を取得する
print(ary1[ary1.nonzero()])
# [ 5  3 -1]

てっきり、array([2, 3, 5]) が戻ってくると予想していたのに、
(array([2, 3, 5]),) というタプルの形で結果が戻ってくる点に注意が必要です。
ただ、上の例の最後に挙げたように、スライスとして使いやすいので便利です。

続いて、2次元配列、要するに行列でやってみましょう。

ary2 = np.array(
        [
            [0, 0, -2, 0],
            [1, 0, 2, 0],
            [0, 0, 1, 0],
            [3, 0, 0, 0],
        ]
    )
print(ary2.nonzero())
# (array([0, 1, 1, 2, 3]), array([2, 0, 2, 2, 0]))

これも直感的には[(0, 2), (1, 0), (1, 2), …]みたいなのが戻ってくると思ったら違う形で帰ってきました。
これは次のようにすると、非ゼロ要素に対する操作を行う場合は次のように使います。(一例です。)

xx, yy = ary2.nonzero()
for x, y in zip(xx, yy):
    print(x, y)
"""
0 2
1 0
1 2
2 2
3 0
"""

また、1次元の場合と同じ様に次のようにしてスライスとして使えます。

print(ary2[ary2.nonzero()])
# [-2  1  2  1  3]