Digamos que tengo un dataframe de Pandas con 4 filas y 5 columnas. Para simplificar, lo convertiré en una matriz Numpy, que se ve así:
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]])Para cada fila, quiero insertar algunas filas similares justo después de la fila, con solo la columna 4 disminuyendo en 1 cada vez hasta 0. El resultado esperado debería verse así:
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]])A continuación se muestra el código que se me ocurrió:
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)Producción
[[ 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]]Obviamente, el código no es eficiente ya que no usa ninguna vectorización Numpy. En realidad, tengo más de 4000 filas originales, por lo que definitivamente se necesita acelerar. Además, no puedo insertar nuevas filas justo después de cada fila. ¿Hay alguna forma eficiente de hacer esto en Numpy o Pandas?
Si está interesado en una solución que no utilice bucles de Python puro sino Numba , aquí tiene una:
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) Esta solución es 12 veces más rápida que la solución de @sammywemmy en mi máquina, aunque A es muy pequeña. Debería ser aún más rápido en entradas más grandes.
arr = A[:, -1] + 1 temp = np.repeat(A, arr, axis = 0) # depending on your array size # you can build the range here with a much faster implementation from this link : # https://stackoverflow.com/a/47126435/7175713 arr = np.concatenate([np.arange(ent) for ent in arr]) temp[:, -1] = temp[:, -1] - arr temp 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]])Finalmente he encontrado una solución:
rep = np.repeat(A, A[:, -1] + 1, axis=0) rep[:, -1] = np.concatenate([np.arange(0, n+1)[::-1] for n in A[:, -1]])Producción:
>>> rep 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]])