时间戳

首页 > 解决方案 > 逐行查找与二维数组条件匹配的元素索引

python - 逐行查找与二维数组条件匹配的元素索引

问题描述

我有一个带有浮点数的二维数组，称为 A，尺寸为 nx m。我有一个名为 B 的维数为 n 的一维浮点数组。对于 AI 中的每一行 i 想要找到小于 B 中第 i 个值的值的索引

我尝试使用 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])]

而且，当然，我希望它被矢量化而不是循环。

标签： pythonnumpyvectorization

解决方案

使用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)]

推荐阅读