python - 如何使用二维数组索引 df 并查找另一个数组中的值？

Question

我的输入是 df 和两个数组:

df = pd.DataFrame({'id': ['id1', 'id2', 'id3', 'id4', 'id5']})

indexes = np.array(
    [[1, 2],
    [4, 0],
    [0, 1],
    [2, 0],
    [1, 0]])

values = np.array(
    [[0.012, 0.019],
    [0.009, 0.012],
    [0.019, 0.028],
    [0.042, 0.061],
    [0.009, 0.021]])

我试图根据indexes数组获取相应的id，同时提取值。

下面的代码给出了预期的输出，但它不仅给了我警告，而且在我的数据集上也非常慢。

wanted = df.copy()
for i, j in enumerate(indexes):
    wanted.at[i, 'list_ids'] = ', '.join(df.iloc[j].squeeze().tolist())
    
for i, j in enumerate(values):
    wanted.at[i, 'list_values'] = np.array(j, dtype='object')
    
print(wanted)

   id   list_ids     list_values
0  id1  id2, id3  [0.012, 0.019]
1  id2  id5, id1  [0.009, 0.012]
2  id3  id1, id2  [0.019, 0.028]
3  id4  id3, id1  [0.042, 0.061]
4  id5  id2, id1  [0.009, 0.021]

你们知道如何改进它或者有其他建议吗？

Answer 1

只需使用 numpy索引:

df['list_ids'] = df['id'].to_numpy()[indexes].tolist()
df['list_values'] = values.tolist()

输出:

    id    list_ids     list_values
0  id1  [id2, id3]  [0.012, 0.019]
1  id2  [id5, id1]  [0.009, 0.012]
2  id3  [id1, id2]  [0.019, 0.028]
3  id4  [id3, id1]  [0.042, 0.061]
4  id5  [id2, id1]  [0.009, 0.021]

如果你想要字符串，不幸的是你必须循环:

df['list_ids'] = list(map(', '.join, df['id'].to_numpy()[indexes]))
df['list_values'] = values.tolist()

输出:

    id  list_ids     list_values
0  id1  id2, id3  [0.012, 0.019]
1  id2  id5, id1  [0.009, 0.012]
2  id3  id1, id2  [0.019, 0.028]
3  id4  id3, id1  [0.042, 0.061]
4  id5  id2, id1  [0.009, 0.021]