How to index a df with a 2D array and lookup for values from another one?

October 25, 2023

My inputs are df and two arrays :

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]])

I’m trying to get the corresponding ids based on the indexes array and also at the same time pull up the values.

My code below gives the expected output but not only it gives me warning but it is also very slow on my dataset.

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]

Do you guys know how to improve it or do you have any other suggestions ?

>Solution :

Simply use numpy indexing:

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

Output:

    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]

If you want strings, unfortunately you have to loop:

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

Output:

    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]