Python中如何批量删除不等长轨迹数组的z坐标？

Got it, let's break down how to solve this problem efficiently. You're dealing with a 1D NumPy array (arr.shape = (3000,)) where each element is a 2D array representing a trajectory with shape (3, N) (x/y/z axes, N time points varying between 150-250). You already know how to modify a single trajectory, but need to scale this to all 3000 entries.

Here are the most straightforward and efficient methods:

Method 1: List Comprehension + Slicing (Recommended) #

Slicing is faster than np.delete because it directly extracts the rows you need instead of creating a new array by removing unwanted rows. This is the simplest and most performant approach for your use case:

# Extract x and y axes (first two rows) from every trajectory
processed_arr = np.array([traj[:2, :] for traj in arr])

• traj[:2, :] grabs the first two rows (x and y) for all time points in the trajectory.
• The result will still be a 1D NumPy array (shape=(3000,)) where each element is a (2, N) array matching the original trajectory's time point count.

Method 2: List Comprehension + np.delete #

If you prefer using np.delete (since you already know it works for single trajectories), you can wrap it in a list comprehension too:

processed_arr = np.array([np.delete(traj, 2, axis=0) for traj in arr])

This does exactly what you'd expect: removes the 3rd row (index 2, z-axis) from each trajectory. Just note that this is slightly less efficient than slicing because np.delete has to handle the deletion logic internally.

Method 3: Using np.vectorize (For "Vectorized" Syntax) #

If you want a more "NumPy-style" vectorized look (even though it's still under-the-hood looping), you can use np.vectorize:

# Define a helper function to remove the z-axis
def strip_z(traj):
    return traj[:2, :]

# Vectorize the function to apply it across the array
vectorized_strip = np.vectorize(strip_z, otypes=[object])
processed_arr = vectorized_strip(arr)

• The otypes=[object] parameter is critical here: it tells NumPy that each output element is an object (in this case, a variable-length array), since we can't fit all trajectories into a uniform 3D array.

Key Note #

Since your trajectories have variable lengths, you can't convert the original array into a single 3D NumPy array (like (3000, 3, 250)). All the methods above will preserve the original 1D object array structure, with each element now being a (2, N) trajectory.

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