使用 NumPy 从另一个数组及其索引创建一个二维数组答案

【问题标题】：Create a 2D array from another array and its indices with NumPy使用 NumPy 从另一个数组及其索引创建一个二维数组
【发布时间】：2017-08-24 09:08:32
【问题描述】：

给定一个数组：

arr = np.array([[1, 3, 7], [4, 9, 8]]); arr

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

并给出它的索引：

np.indices(arr.shape)

array([[[0, 0, 0],
        [1, 1, 1]],

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

我怎样才能将它们整齐地堆叠在一起以形成一个新的二维阵列？这就是我想要的：

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

这是我目前的解决方案：

def foo(arr):
    return np.hstack((np.indices(arr.shape).reshape(2, arr.size).T, arr.reshape(-1, 1)))

它有效，但是执行此操作是否有更短/更优雅的方法？

【问题讨论】：

如果数组是与 np.intp 不同的数据类型会发生什么？输出应该是什么类型？
@Eric 啊，我明白你的意思了。如果数组是浮点数，我认为可以将索引转换为浮点数。

标签： python arrays numpy

【解决方案1】：

使用array-initialization 然后broadcasted-assignment 在后续步骤中分配索引和数组值 -

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

请注意，我们正在避免使用 np.indices(arr.shape)，这可能会减慢速度。

示例运行 -

In [10]: arr = np.array([[1, 3, 7], [4, 9, 8]])

In [11]: indices_merged_arr(arr)
Out[11]: 
array([[0, 0, 1],
       [0, 1, 3],
       [0, 2, 7],
       [1, 0, 4],
       [1, 1, 9],
       [1, 2, 8]])

性能

arr = np.random.randn(100000, 2)

%timeit df = pd.DataFrame(np.hstack((np.indices(arr.shape).reshape(2, arr.size).T,\
                                arr.reshape(-1, 1))), columns=['x', 'y', 'value'])
100 loops, best of 3: 4.97 ms per loop

%timeit pd.DataFrame(indices_merged_arr_divakar(arr), columns=['x', 'y', 'value'])
100 loops, best of 3: 3.82 ms per loop

%timeit pd.DataFrame(indices_merged_arr_eric(arr), columns=['x', 'y', 'value'], dtype=np.float32)
100 loops, best of 3: 5.59 ms per loop

注意： 时间包括转换为 pandas 数据帧，这是此解决方案的最终用例。

【讨论】：

好吧，这看起来很简单。您是否会考虑为更大的 2D 数组添加一些计时，只是为了完整性？
@cᴏʟᴅsᴘᴇᴇᴅ 你有我可以比较的循环解决方案吗？
我已经编辑了我在问题中的解决方案，如果有帮助的话。这是我唯一的解决方案。
@cᴏʟᴅsᴘᴇᴇᴅ 似乎没有任何好转。我想一个更好的可以使用this one。
添加了一些性能统计信息。您的解决方案很棒！

【解决方案2】：

对于 nd 数组的更通用的答案，可以正确处理其他 dtypes：

def indices_merged_arr(arr):
    out = np.empty(arr.shape, dtype=[
        ('index', np.intp, arr.ndim),
        ('value', arr.dtype)
    ])
    out['value'] = arr
    for i, l in enumerate(arr.shape):
        shape = (1,)*i + (-1,) + (1,)*(arr.ndim-1-i)
        out['index'][..., i] = np.arange(l).reshape(shape)
    return out.ravel()

这将返回一个结构化数组，其中包含一个索引列和一个值列，它们可以是不同的类型。

【讨论】：