使用 NumPy 从另一个数组及其索引创建一个二维数组 [英] Create a 2D array from another array and its indices with NumPy

问题描述

给定一个数组:

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]]])

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

How would I be able to stack them neatly one against the other to form a new 2D array? This is what I'd like:

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

这是我目前的解决方案:

This is my current solution:

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

它有效，但有没有更短/更优雅的东西来执行这个操作?

It works, but is there something shorter/more elegant to carry this operation out?

推荐答案

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

Using array-initialization and then broadcasted-assignment for assigning indices and the array values in subsequent steps -

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)，这可能会减慢速度.

Note that we are avoiding the use of np.indices(arr.shape), which could have slowed things down.

样品运行 -

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 数据帧，这是此解决方案的最终用例.

Note: Timings include conversion to pandas dataframe, that is the eventual use case for this solution.

这篇关于使用 NumPy 从另一个数组及其索引创建一个二维数组的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！