使用索引展平一个numpy数组

时间:2020-06-13 18:09:22

标签: python arrays numpy

给出一个形状为(X, Y, 2)的NumPy数组,该数组表示包括“点”的“框架”的数组,每个点的坐标为(x,y),我想将第一维和第二维合并为一个(X*Y, 4)数组现在代表所有点以及X和Y维度的索引。

例如,如果我的数组是:

[
  [          # Frame 0
   [1, 2],   # Point 0
   [2, 3]    # Point 1
  ],
  [          # Frame 1
   [4, 5],   # Point 0
   [6, 7]    # Point 1
  ]
]

我想获取数组:

[
  [0, 0, 1, 2],   # Frame 0, Point 0
  [0, 1, 2, 3]    # Frame 0, Point 1
  [1, 0, 4, 5],   # Frame 1, Point 0
  [1, 1, 6, 7]    # Frame 1, Point 1
]

慢速解决方案:

arr = np.array([[[1, 2],[2, 3]],[[4, 5],[6, 7]]])
new_arr = []
for i, points in enumerate(arr):
  for j, point in enumerate(points):
    new_arr.append([i, j] + point.tolist())

有更快的方法吗?

3 个答案:

答案 0 :(得分:2)

您可以使用numpy.ndindex分别求解每个部分,以获取索引和.reshape()。然后,您可以使用numpy.c_来堆叠它们。

a = np.array([[[1, 2],[2, 3]],[[4, 5],[6, 7]]])
c = a.reshape(-1, a.shape[-1])
print(c)
# [[1 2]
#  [2 3]
#  [4 5]
#  [6 7]]
indices = list(np.ndindex(a.shape[:-1]))
print(indices)
# [(0, 0), (0, 1), (1, 0), (1, 1)]
print(np.c_[indices, c])
# [[0 0 1 2]
#  [0 1 2 3]
#  [1 0 4 5]
#  [1 1 6 7]]

答案 1 :(得分:1)

此代码中使用了一个更大的示例数组,以便可以在每个维度上以不同的大小对其进行测试:

import numpy as np

arr = np.array(
    [
        [          
            [1, 2],   
            [2, 3],   
            [3, 4]    
            ],
        [          
            [4, 5],   
            [6, 7],    
            [8, 7]    
            ],
        [          
            [14, 5],   
            [16, 7],    
            [18, 7]    
            ],
        [            
            [24, 5],   
            [26, 7],    
            [28, 7]    
            ]
        ]
)

x, y = arr.shape[:2]
assert(arr.shape[2] == 2)
ay, ax = (a.reshape(x, y, 1) for a in np.meshgrid(np.arange(y), np.arange(x)))
new_array = np.concatenate([ax, ay, arr], axis=2).reshape(x * y, 4)

print(repr(new_array))

给出以下内容:

array([[ 0,  0,  1,  2],
       [ 0,  1,  2,  3],
       [ 0,  2,  3,  4],
       [ 1,  0,  4,  5],
       [ 1,  1,  6,  7],
       [ 1,  2,  8,  7],
       [ 2,  0, 14,  5],
       [ 2,  1, 16,  7],
       [ 2,  2, 18,  7],
       [ 3,  0, 24,  5],
       [ 3,  1, 26,  7],
       [ 3,  2, 28,  7]])

使用原始示例数组可以得出:

array([[0, 0, 1, 2],
       [0, 1, 2, 3],
       [1, 0, 4, 5],
       [1, 1, 6, 7]])

没有显式循环,因此它应该更快。 (任何循环都在numpy内部,并将以优化的C代码实现。)

答案 2 :(得分:0)

我也是NumPy的新手,但我认为这应该可以工作(如果我错了,请纠正我):arr.reshape(-1,4)