python - 在 Numpy 或 Pandas 中的每一行之后插入相似行(只有一列发生变化)的有效方法

Question

假设我有一个包含 4 行和 5 列的 Pandas 数据框。为简单起见，我会将其转换为 Numpy 数组，如下所示:

import numpy as np

A = np.array([[23, 43, 23, 110, 5],
              [83, 32, 12, 123, 4],
              [58, 41, 59, 189, 1],
              [93, 77, 22, 170, 3]])

对于每一行，我想在该行之后插入一些相似的行，只有第 4 列每次减 1 直到 0。预期输出应该如下所示:

np.array([[23, 43, 23, 110, 5],
          [23, 43, 23, 110, 4],
          [23, 43, 23, 110, 3],
          [23, 43, 23, 110, 2],
          [23, 43, 23, 110, 1],
          [23, 43, 23, 110, 0],

          [83, 32, 12, 123, 4],
          [83, 32, 12, 123, 3],
          [83, 32, 12, 123, 2],
          [83, 32, 12, 123, 1],
          [83, 32, 12, 123, 0],

          [58, 41, 59, 189, 1],
          [58, 41, 59, 189, 0],

          [93, 77, 22, 170, 3],
          [93, 77, 22, 170, 2],
          [93, 77, 22, 170, 1],
          [93, 77, 22, 170, 0]])

下面是我想出的代码:

new_rows = []
for i, row in enumerate(A):
    new = A[i, 4] - 1
    while new >= 0:
        new_row = row.copy()
        new_row[4] = new
        new_rows.append(new_row)
        new -= 1
new_A = np.vstack([A, np.array(new_rows)])
print(new_A)

输出

[[ 23  43  23 110   5]
 [ 83  32  12 123   4]
 [ 58  41  59 189   1]
 [ 93  77  22 170   3]
 [ 23  43  23 110   4]
 [ 23  43  23 110   3]
 [ 23  43  23 110   2]
 [ 23  43  23 110   1]
 [ 23  43  23 110   0]
 [ 83  32  12 123   3]
 [ 83  32  12 123   2]
 [ 83  32  12 123   1]
 [ 83  32  12 123   0]
 [ 58  41  59 189   0]
 [ 93  77  22 170   2]
 [ 93  77  22 170   1]
 [ 93  77  22 170   0]]

显然，代码效率不高，因为它没有使用任何 Numpy 向量化。实际上，我有 4,000 多行原始行，因此绝对需要加速。此外，我不能在每一行之后立即插入新行。在 Numpy 或 Pandas 中有什么有效的方法可以做到这一点吗？

Answer 1

如果您对不使用任何纯 Python 循环而是使用 Numba 的解决方案感兴趣，那么这里有一个:

import numba as nb

# 2 overloads: one for np.int32 types and one for np.int64 regarding the type of A
@nb.njit(['int32[:,::1](int32[:,::1])', 'int64[:,::1](int64[:,::1])'])
def compute(A):
    n, m = A.shape
    rows = A[:,-1].sum() + n
    res = np.empty((rows, m), dtype=A.dtype)
    row = 0
    for i in range(n):
        count = A[i, -1]
        for j in range(count+1):
            res[row+j, 0:m-1] = A[i, 0:m-1]
            res[row+j, m-1] = count-j
        row += count+1
    return res

result = compute(A)

尽管 A 非常小，但此解决方案比我机器上@sammywemmy 的解决方案快 12 倍。它应该在更大的输入上更快。