按行查找与2D数组条件匹配的元素的索引

时间:2019-07-04 14:56:41

标签: python numpy vectorization

我有一个2D数组,其浮点称为A,尺寸为n x m。我有一个维数为n的浮点数的一维数组,称为B。对于A中的每一行,我都想找到小于B中第ith个值的值的索引。

我尝试使用np.argwhere

A = np.array([[.1,.2,.3,.4], [.1,.2,.3,.4], [.1,.2,.3,.4]]) B = np.array([0.15, 0.25, 0.35]) np.argwhere(A.T < B) 这是输出

array([[0, 0],[0, 1],[0, 2],[1, 1],[1, 2],[2, 2]])

我想拥有这样的东西

[array([0,1,2]), array([1,2]), array([2])]

当然,我希望将其向量化而不循环播放。

1 个答案:

答案 0 :(得分:0)

使用np.where + np.split

x, y = np.where(A.T < B)

s = np.flatnonzero(np.append([False], x[1:] != x[:-1]))

np.split(y, s)

[array([0, 1, 2], dtype=int64),
 array([1, 2], dtype=int64),
 array([2], dtype=int64)]