代码之家 › 专栏 › 技术社区 › Amit

使用索引展平numpy数组

numpy arrays python

Amit · 技术社区 · 6 年前

给定一个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())

有没有更快的办法?

2 回复 | 直到 6 年前

V. Ayrat 6 年前

您可以使用 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]]

alani 6 年前

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

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

theX 6 年前