用索引扁平化一个numpy数组。

Question

给定形状的NumPy数组 (X, Y, 2) 代表包括 "点 "在内的 "帧 "数组，每个点都有一个 "帧"。(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())

有没有更快的方法？

Answer 1

你可以用以下方法分别解决每个部分 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]]

Answer 2

在这段代码中使用了一个较大的例子数组，这样就可以用每个维度的不同大小来测试。

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代码中实现）。

Answer 3

我也是NumPy的新手，但我认为这样做应该可以（如果我错了，有人纠正我）：arr.reshape(-1,4)