根据索引数组连接 DataFrame 中的值

Question

我有一个 DataFrame test，形状为 (1138812, 57)。头长这样：

我有一个数组 indices，其形状为 (1138812, 25)。它是一个二维数组，每个子数组有 25 个索引。它看起来像这样：

[

indices 数组有来自 DataFrame 的 25 个索引对应于来自同一 DataFrame 的每个 1138812 个索引。我想基于此数组创建一个包含 25 X 1138812 行的新 DataFrame。例如：

我有一个二维数组，类似于：

[[0,2,3],
 [1,0,3],
 [2,1,0],
 [3,1,2]]

我有一个 pandas 数据框，类似于：

现在我想根据二维数组中列出的索引获取一个新的数据框，对于这个例子，它就像：

 id   val   id_2  val
0 a    9      a    9
0 a    9      c    3
0 a    9      d    7
1 b    8      b    8
1 b    8      a    9
1 b    9      d    7
2 c    3      c    3
2 c    3      b    8
2 c    3      a    9
3 d    7      d    7
3 d    7      b    8
3 d    7      c    3

我尝试了很多方法，包括：

import pandas as pd
import numpy as np

index = [[0,2,3],
 [1,0,3],
 [2,1,0],
 [3,1,2]]

idse = ['a','b','c','d']
vals = [9,8,3,7]

data = {'id': idse, 'val': vals}
df = pd.DataFrame(data=data)

newdf = pd.DataFrame(np.repeat(df.values, len(index[0]), axis=0))

flat_list = [item for sublist in index for item in sublist]
newdf['id_2'] = df.id[flat_list].values
newdf['val_2'] = df.val[flat_list].values

和

fdf = pd.DataFrame()
for i, ir in enumerate(l):
  temp_df = df.iloc[ir]
  temp_df['id'] = df.iloc[i]['id']
  temp_df = pd.merge(df,temp_df,how="outer",on="id")
  temp_df = temp_df[temp_df['id']==df.iloc[i]['id']]
  fdf = pd.concat([fdf,temp_df])
fdf

它们都按照我想要的方式工作，但对于具有 1.1m 行的原始 DataFrame 来说它们非常非常慢，并且它们占用了所有导致笔记本崩溃的内存。我正在使用 RAPIDS 库，包括分别对应于 pandas、numpy/scipy 和 sklearn 的 cuDF、cuPy、cuML，我需要一个纯粹的 numpy/pandas 解决方案，以便它们可以使用 GPU 内核并制作它操作更快更高效

谢谢

Answer 1

假设 df 和 a 输入数据帧和数组，您可以 repeat 数据帧的索引和 concat 它与从展平数组索引的数据帧:

idx = df.index.repeat(a.shape[1])
df2 = pd.concat(
          [df.loc[idx],
           df.loc[a.ravel()].add_suffix('_2').set_axis(idx)
          ], axis=1)

输出：

  id  val id_2  val_2
0  a    9    a      9
0  a    9    c      3
0  a    9    d      7
1  b    8    b      8
1  b    8    a      9
1  b    8    d      7
2  c    3    c      3
2  c    3    b      8
2  c    3    a      9
3  d    7    d      7
3  d    7    b      8
3  d    7    c      3

使用的输入：

df = pd.DataFrame({'id': ['a', 'b', 'c', 'd'],
                   'val': [9, 8, 3, 7]})

a = np.array([[0,2,3],
              [1,0,3],
              [2,1,0],
              [3,1,2]])

注意。快速测试表明处理 1M 行需要 900 毫秒

Answer 2

让我们尝试 assign 和 explode 然后 join

out = df.assign(new=a.tolist()).explode('new').set_index('new').add_suffix('_2').join(df)

根据索引数组连接 DataFrame 中的值

Join values from a DataFrame according to an array of indices

arrays

numpy

pandas

rapids

cudf