python - 获取二维numpy数组中大于阈值的元素的列数

Question

我有一个像这样的数组，希望返回值超过阈值 0.6 的每一行的列号:

X = array([[ 0.16,  0.40,  0.61,  0.48,  0.20],
        [ 0.42,  0.79,  0.64,  0.54,  0.52],
        [ 0.64,  0.64,  0.24,  0.63,  0.43],
        [ 0.33,  0.54,  0.61,  0.43,  0.29],
        [ 0.25,  0.56,  0.42,  0.69,  0.62]])

结果将是:

[[2],
[1, 2],
[0, 1, 3],
[2],
[3, 4]]

有没有比双 for 循环更好的方法？

def get_column_over_threshold(data, threshold):
    column_numbers = [[] for x in xrange(0,len(data))]
    for sample in data:
        for i, value in enumerate(data):
            if value >= threshold:
                column_numbers[i].extend(i)
    return topic_predictions

Answer 1

对于每一行，您可以询问元素大于 0.6 的索引:

result = [where(row > 0.6) for row in X]

这会执行您想要的计算，但是 result 的格式有点不方便，因为在这种情况下 where 的结果是一个 tuple 大小为 1，包含带有索引的 NumPy 数组。我们可以将 where 替换为 flatnonzero 来直接获取数组而不是元组。为了获取列表的列表，我们显式地将此数组转换为列表:

result = [list(flatnonzero(row > 0.6)) for row in X]

(在上面的代码中，我假设您使用了 from numpy import *)