Extract n columns with highest sum in NumPy array

Question

Imagine I have a NumPy matrix with 100 rows and 1000 columns.

How do I get a new matrix composed by the n columns that have the highest sums in the original matrix?

Answer 1

You can use np.argsort as done by @NPE here. Here's an example on two smaller arrays:

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]]