任意维 np 矩阵输出满足条件的值 位置

统计数据中满足某个条件的数的个数，并将其位置输出

首先看 二维 np 矩阵

import numpy as np

a = np.arange(1, 16).reshape((3, 5))
print(a)
b = np.where(a % 2 == 0)
print(b)

[[ 1  2  3  4  5]
 [ 6  7  8  9 10]
 [11 12 13 14 15]]
(array([0, 0, 1, 1, 1, 2, 2], dtype=int64), array([1, 3, 0, 2, 4, 1, 3], dtype=int64))

可以看到 np.where 返回的结果中，

• 第1个array表示满足条件的数字的行号
• 第2个array表示满足条件的数字的列号

高维矩阵依次类推，可以得到任意维矩阵找出满足条件下标的一般算法：

def get_cond_idx(X, cond=None):
    find = np.where(cond)  # condition: bool array like
    idx = find[0]  # first dimension, at least has one array find[0]
    for i in range(1, X.ndim):
        idx = np.vstack((idx, find[i]))
    return {
        'data': X[find],
        'idx': idx.T,
        'num': find[0].__len__()
    }

# 二维数据
a = np.arange(1, 16).reshape((3, 5))
print(a)
b = get_cond_idx(a, cond=a % 2 == 0)
for i in range(b['num']):
    print(i, b['idx'][i], b['data'][i])

[[ 1  2  3  4  5]
 [ 6  7  8  9 10]
 [11 12 13 14 15]]
0 [0 1] 2
1 [0 3] 4
2 [1 0] 6
3 [1 2] 8
4 [1 4] 10
5 [2 1] 12
6 [2 3] 14

# 三维数据
a = np.arange(1, 25).reshape((2, 3, 4))
print(a)
b = get_cond_idx(a, cond=a % 5 == 0)
for i in range(b['num']):
    print(i, b['idx'][i], b['data'][i])

[[[ 1  2  3  4]
  [ 5  6  7  8]
  [ 9 10 11 12]]

 [[13 14 15 16]
  [17 18 19 20]
  [21 22 23 24]]]
0 [0 1 0] 5
1 [0 2 1] 10
2 [1 0 2] 15
3 [1 1 3] 20