Pandas dataframe - Column contains index to other columns

I have a dataframe (df) where one of the columns ('bestcol') contains indexes of other columns in the table. I want to grab the column being referred to by 'bestcol', round it, and create a new column with that info (see table below for rough example, in which bestcol = 1 refers to Val1, 2 refers to Val2, 3 refers to Val3).

My initial approach involved looping through the table row-by-row:

bestcol_list = []
for i in range(len(df)):
    bestVal = round(df.iloc[i, df['bestcol'][i]], 0)
    bestcol_list.append(bestVal)
    
df['Final'] = bestcol_list

My data has a few million records, so this was a time consuming process. My next approach involved using apply:

bestcol_list = df.apply(lambda row: round(row[row['bestcol']], 0), axis=1)

df['Final'] = bestcol_list

This turned out to be a bit slower than just looping through the table. Is there a vectorized approach to solving this problem that I'm not considering?

Thanks!

You can use numpy indexing:

row = np.arange(len(df))
col = df['bestcol'].values - 1
x = df.filter(like='Val').values  # or df.iloc[:, :3].values

df['Final'] = np.round(x[row, col])

Output:

>>> df
    Val1   Val2   Val3  bestcol  Final
0    1.1    2.1    3.1        1    1.0
1   11.1   22.1   33.1        2   22.0
2  111.1  222.1  333.1        3  333.0

Performance for 5_000_000 rows and 100 columns:

M = 5_000_000
N = 100
x = np.random.uniform(1, 500, (M, N))
row = np.arange(M)
col = np.random.randint(1, N+1, M) - 1

%timeit np.round(x[row, col])
75.5 ms ± 388 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

I used nested np.where to handle three conditions and then a np.round.

df['Final'] = np.round(np.where(df.bestcol == 1, df.Val1, np.where(df.bestcol == 2, df.Val2, df.Val3)), 0)

A simple and fast pythonic way :

import pandas as pd

df = pd.DataFrame({'Val1': [1.1, 11.1, 111.1],
                   'Val2': [2.1, 22.1, 222.1],
                   'Val3': [3.1, 33.1, 333.1],
                   'bestcol': [1, 2, 3],
                   })

df['Final'] = [round(df[['Val1', 'Val2', 'Val3']].iloc[i, p]) \
               for i, p in enumerate(df.bestcol.sub(1).tolist())]

print(df)

Result

    Val1   Val2   Val3  bestcol  Final
0    1.1    2.1    3.1        1      1
1   11.1   22.1   33.1        2     22
2  111.1  222.1  333.1        3    333

Since your bestcol contains an ordered ordinal positions of actual columns you can apply numpy.diag:

df['Final'] = np.round(np.diag(df[df.columns[:-1]]))

       1      2      3  bestcol  Final
0    1.1    2.1    3.1        1    1.0
1   11.1   22.1   33.1        2   22.0
2  111.1  222.1  333.1        3  333.0

try this:

df['Final'] = df.values[df.index, df.bestcol-1].round(0)
print(df)
>>>
    Val1   Val2   Val3  bestcol  Final
0    1.1    2.1    3.1        1    1.0
1   11.1   22.1   33.1        2   22.0
2  111.1  222.1  333.1        3  333.0

</details>