如何更简洁地对NumPy数组分组列求和并重组数组？

Great question! For this kind of column-wise aggregation and combining, using np.c_ (a shorthand for column concatenation) makes the code much cleaner and more readable than np.concatenate. Here's a concise solution:

First, let's set up our original array:

import numpy as np
a = np.arange(21).reshape(3,7)

Now compute the desired result in one clean step:

result = np.c_[
    a[:, :2],                  # Keep the first 2 columns as-is
    a[:, 2:5].sum(axis=1),     # Sum columns 2, 3, 4 (indices 2 to 4 inclusive)
    a[:, 5:].sum(axis=1)       # Sum columns 5 and 6
]

Running this gives exactly the output you want:

print(result)
# Output:
# array([[ 0,  1,  9, 11],
#        [ 7,  8, 30, 25],
#        [14, 15, 51, 39]])

Why this works: #

• np.c_ automatically handles stacking 1D arrays (like the sums from sum(axis=1)) as columns alongside 2D arrays, so we don't need extra code to maintain a 2D shape (though using keepdims=True would also work if you prefer explicit control).
• The slicing is intuitive: a[:, :2] grabs the first two columns, a[:,2:5] targets columns 2-4 (since slicing stops just before the end index), and a[:,5:] takes everything from column 5 onwards.

If you prefer a more explicit function name, np.column_stack works identically here:

result = np.column_stack([
    a[:, :2],
    a[:,2:5].sum(axis=1),
    a[:,5:].sum(axis=1)
])

Both methods are more concise than np.concatenate because they eliminate the need to specify axis=1 and handle shape alignment gracefully.

内容的提问来源于stack exchange，提问作者Keyur Potdar