小能豆

在 NumPy 数组中概括切片操作

py

这个问题基于这个老问题:

给定一个数组:

py In [122]: arr = np.array([[1, 3, 7], [4, 9, 8]]); arr Out[122]: array([[1, 3, 7], [4, 9, 8]])

并给出其指数:

```py
In [127]: np.indices(arr.shape)
Out[127]:
array([[[0, 0, 0],
[1, 1, 1]],

   [[0, 1, 2],
    [0, 1, 2]]])

```

我如何才能将它们整齐地堆叠在一起以形成一个新的 2D 数组?这就是我想要的:

py array([[0, 0, 1], [0, 1, 3], [0, 2, 7], [1, 0, 4], [1, 1, 9], [1, 2, 8]])

Divakar 提出的这个解决方案是我目前用于二维数组的解决方案:

def indices_merged_arr(arr):
    m,n = arr.shape
    I,J = np.ogrid[:m,:n]
    out = np.empty((m,n,3), dtype=arr.dtype)
    out[...,0] = I
    out[...,1] = J
    out[...,2] = arr
    out.shape = (-1,3)
    return out

现在,如果我想传递一个 3D 数组,我需要修改此函数:

def indices_merged_arr(arr):
    m,n,k = arr.shape   # here
    I,J,K = np.ogrid[:m,:n,:k]   # here
    out = np.empty((m,n,k,4), dtype=arr.dtype)   # here
    out[...,0] = I
    out[...,1] = J
    out[...,2] = K     # here
    out[...,3] = arr
    out.shape = (-1,4)   # here
    return out

但是此函数现在仅适用于 3D 数组 - 我无法将 2D 数组传递给它。

有没有办法可以将其推广到任何维度?这是我的尝试:

def indices_merged_arr_general(arr):
    tup = arr.shape   
    idx = np.ogrid[????]   # not sure what to do here....
    out = np.empty(tup + (len(tup) + 1, ), dtype=arr.dtype) 
    for i, j in enumerate(idx):
        out[...,i] = j
    out[...,len(tup) - 1] = arr
    out.shape = (-1, len(tup)
    return out

我对这一行感到困惑:

idx = np.ogrid[????]   

我怎样才能让它工作?


阅读 21

收藏
2024-10-10

共1个答案

小能豆

您可以创建该函数的通用版本,indices_merged_arr通过根据输入数组的形状动态构建索引数组来处理任意维度的数组。操作方法如下:

  1. 用于np.indices获取输入数组维度的索引网格。
  2. 将索引和原始数组堆叠在一起以形成所需的输出形状。
  3. 使用np.concatenate来正确地结合它们。

完整的实现如下:

import numpy as np

def indices_merged_arr_general(arr):
    # Get the shape of the input array
    shape = arr.shape
    # Generate a grid of indices for all dimensions
    idx = np.indices(shape)

    # Stack the indices and the array along a new last axis
    out = np.concatenate([idx.reshape(len(shape), -1).T, arr.flatten()[:, np.newaxis]], axis=1)

    return out

# Example usage:
arr_2d = np.array([[1, 3, 7], [4, 9, 8]])
print(indices_merged_arr_general(arr_2d))

arr_3d = np.random.rand(2, 3, 4)  # Example 3D array
print(indices_merged_arr_general(arr_3d))

解释:

  1. np.indices(shape)为数组的所有维度生成一个索引网格,返回一个形状为 的数组(d, *shape),其中d是维数。
  2. idx.reshape(len(shape), -1).T将索引重塑为二维数组,其中每一行对应于原始数组中每个元素的索引元组。这-1允许它根据输入形状自动计算元素数量。
  3. arr.flatten()[:, np.newaxis]将输入数组展平为一维数组,并将其重塑为列向量,以便与索引连接。
  4. np.concatenate([...], axis=1)沿着最后一个轴结合索引和原始值。

输出

对于 2D 示例,输出将是:

array([[0, 0, 1],
       [0, 1, 3],
       [0, 2, 7],
       [1, 0, 4],
       [1, 1, 9],
       [1, 2, 8]])

对于 3D 示例,您将获得一个数组,该数组具有与原始数组深度相同的附加维度。此广义函数适用于任意数量的维度,可有效创建将索引与值相结合的输出

2024-10-10