如何在for循环中使用NumPy正确拼接数组？

Hey there! Let's break down why your current code isn't working and how to fix it to get the array concatenation you're expecting.

What's Wrong with Your Original Code? #

Looking at your code:

import numpy as np
t = np.arange(3)
for i in range(5):
    a = np.arange(3)
    np.stack((t, a))
print(t)

There are two key issues here:

1. You're not saving the result of np.stack: The np.stack function returns a new array instead of modifying input arrays in-place. Since you don't assign this returned value back to t (or any variable), each stack operation's result is discarded immediately.
2. np.stack isn't the best fit for incremental appending: Even if you saved the result, stacking two 1D arrays (t and a) would create a 2D array of shape (2,3) each time. Repeating this in a loop would keep adding dimensions (e.g., next iteration would produce (2,2,3)), which isn't likely what you want if you're trying to build a single 2D array of all your stacked arrays.

Solutions to Fix It #

Option 1: Collect Arrays in a List First (Recommended)

The most efficient way to build a larger array from multiple smaller ones is to collect all arrays in a Python list first, then stack them all at once. This avoids repeated memory reallocations that come with modifying NumPy arrays in a loop.

import numpy as np

# Start with your initial array
arr_list = [np.arange(3)]

# Collect each new array in the list
for i in range(5):
    a = np.arange(3)
    arr_list.append(a)

# Stack all arrays at once along the 0th axis (rows)
t = np.stack(arr_list, axis=0)
print(t)

This outputs a 2D array of shape (6,3) (your initial t plus 5 copies of a). If you don't want the initial t included, just start with an empty list: arr_list = [].

Option 2: Incremental Concatenation in the Loop

If you need to build the array incrementally during the loop, use np.concatenate instead of np.stack, and make sure to reshape your arrays to match dimensions before concatenating:

import numpy as np

# Start with a 2D array (add a new axis to the initial 1D array)
t = np.arange(3)[np.newaxis, :]

for i in range(5):
    # Reshape the new array to 2D as well
    a = np.arange(3)[np.newaxis, :]
    # Concatenate along the 0th axis (rows) and assign back to t
    t = np.concatenate((t, a), axis=0)

print(t)

This also produces a (6,3) array, building it step by step in the loop.

Key Takeaways #

• Always assign the result of NumPy array operations back to a variable if you want to keep the changes.
• Prefer collecting arrays in a list first and stacking once over incremental modifications—this is faster and more efficient for most cases.
• Use np.concatenate for appending arrays along an existing axis, and np.stack when you want to add a new axis to the combined array.

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