过滤numpy数组的行?

时间:2014-10-02 05:00:53

标签: python numpy filter

我希望将函数应用于numpy数组的每一行。如果此函数的计算结果为true,我将保留该行,否则我将丢弃该行。例如,我的功能可能是:

def f(row):
    if sum(row)>10: return True
    else: return False

我想知道是否有类似的东西:

np.apply_over_axes()

将函数应用于numpy数组的每一行并返回结果。我希望有类似的东西:

np.filter_over_axes()

将函数应用于numpy数组的每一行,并仅返回函数返回true的行。有这样的事吗?或者我应该只使用for循环?

1 个答案:

答案 0 :(得分:29)

理想情况下,您可以实现函数的矢量化版本并使用它来执行boolean indexing。对于绝大多数问题,这是正确的解决方案。 Numpy提供了许多可以在各种轴上执行的函数以及所有基本操作和比较,因此大多数有用的条件都应该是可矢量化的。

import numpy as np

x = np.random.randn(20, 3)
x_new = x[np.sum(x, axis=1) > .5]

如果您完全确定不能执行上述操作,我建议您使用列表推导(或np.apply_along_axis)创建一个bool数组来索引。

def myfunc(row):
    return sum(row) > .5

bool_arr = np.array([myfunc(row) for row in x])
x_new = x[bool_arr]

这将以相对干净的方式完成工作,但将比矢量化版本慢得多。一个例子:

x = np.random.randn(5000, 200)

%timeit x[np.sum(x, axis=1) > .5]
# 100 loops, best of 3: 5.71 ms per loop

%timeit x[np.array([myfunc(row) for row in x])]
# 1 loops, best of 3: 217 ms per loop