根据子元素的值重复 NumPy 数组中的元素

Question

我需要重复以下 numpy 数组 arr 中的元素，以便根据子元素的值重复元素。在执行重复时，我想将元素的第三个值重置为 1。

arr = np.array([[A, B, 1],
[B, C, 2],
[C, D, 3],
[D, E, 1]])

所需输出：

[[A, B, 1],
[B, C, 1],
[B, C, 1],
[C, D, 1],
[C, D, 1],
[C, D, 1],
[D, E, 1]]

下面的代码运行良好，但是对于大型数据集（即〜500k长度）来说它变得非常慢。

_arr = primary_auto.to_records()
pm_arr = np.empty((0,3))

for i in range(len(_arr)):
    for nr in range(_arr[i][2]):
        pm_arr = np.append(pm_arr, np.array([[_arr[i][0],_arr[i][1],1]]), axis = 0)

Answer 1

使用

numpy.repeat

，然后分配

1

：

arr = np.array([['A', 'B', 1],
['B', 'C', 2],
['C', 'D', 3],
['D', 'E', 1]])

out = np.repeat(arr, arr[:, 2].astype(int), axis=0)
out[:, -1] = 1

注意，由于是字符串，所有数组都是字符串，需要转换为整数

repeat

输出：

array([['A', 'B', '1'],
       ['B', 'C', '1'],
       ['B', 'C', '1'],
       ['C', 'D', '1'],
       ['C', 'D', '1'],
       ['C', 'D', '1'],
       ['D', 'E', '1']], dtype='<U11')

Answer 2

我通过使用 numpy 找到答案

import numpy as np
arr = np.array([["A", "B", 1],
["B", "C", 2],
["C", "D", 3],
["D", "E", 2]])
k = np.repeat(arr, arr[:, -1].astype(np.int32), 0)
k[:, -1] = 1
k

array([['A', 'B', '1'],
       ['B', 'C', '1'],
       ['B', 'C', '1'],
       ['C', 'D', '1'],
       ['C', 'D', '1'],
       ['C', 'D', '1'],
       ['D', 'E', '1'],
       ['D', 'E', '1']], dtype='<U11')