在NumPy数组中提取具有最高和的n列

时间:2017-10-09 19:52:51

标签: python numpy

想象一下,我有一个包含100行和1000列的NumPy矩阵。

如何获得由原始矩阵中具有最高总和的 n 列组成的新矩阵?

1 个答案:

答案 0 :(得分:0)

您可以像@NPE here一样使用np.argsort。这是两个较小阵列的示例:

def nlargest_cols(a, n):
    return a[:, sorted(a.sum(axis=0).argsort()[-n:][::-1])]

# `a` is a 3x4 array with column sums getting
#     larger from left to right.
a = np.arange(12).reshape(3,4)

# `b` is `a` rotated 2 turns.
b = np.rot90(a, 2)

print(nlargest_cols(a, 2))
# [[ 2  3]
#  [ 6  7]
#  [10 11]]

print(nlargest_cols(b, 3))
# [[11 10  9]
#  [ 7  6  5]
#  [ 3  2  1]]