在每个元素上使用条件的 Numpy 过滤器

Question

我有一个过滤器表达式如下：

feasible_agents = filter(lambda agent: agent >= cost[task, agent], agents)

agentspython 列表在哪里。

现在，为了加快速度，我正在尝试使用 numpy 来实现这一点。

使用 numpy 的等价物是什么？

我知道这有效：

threshold = 5.0
feasible_agents = np_agents[np_agents > threshold]

np_agents的 numpy 等价物在哪里agents？

但是，我希望阈值是 numpy 数组中每个元素的函数。

Answer 1

您可以使用numpy.extract：

>>> nparr = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
>>> nparreven = np.extract(nparr % 2 == 0, nparr)

或numpy.where：

>>> nparr = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
>>> nparreven = nparr[np.where(nparr % 2 == 0)]

Answer 2

由于您没有提供示例数据，因此使用玩具数据：

# Cost of agents represented by indices of cost, we have agents 0, 1, 2, 3
cost = np.array([4,5,6,2])
# Agents to consider 
np_agents = np.array([0,1,3])
# threshold for each agent. Calculate different thresholds for different agents. Use array of indexes np_agents into cost array.
thresholds = cost[np_agents] # np.array([4,5,2])
feasible_agents = np_agents[np_agents > thresholds] # np.array([3])