I am new to Python/Pandas. Consider the following code:
import pandas as pd
import numpy as np
df = pd.DataFrame({'Time': [0.0, 1.0, 2.0, 0.0, 1.0, 2.0, 0.0, 2.0, 0.0, 1.0, 2.0],
'Id': [1, 1, 1, 2, 2, 2, 3, 3, 4, 4, 4],
'A': [10, 15, np.NaN, 11, 16, 25, 10, 15, 9, 14, 19]})
print(df)
Output:
A Id Time
0 10.0 1 0.0
1 15.0 1 1.0
2 NaN 1 2.0
3 11.0 2 0.0
4 16.0 2 1.0
5 25.0 2 2.0
6 10.0 3 0.0
7 15.0 3 2.0
8 9.0 4 0.0
9 14.0 4 1.0
10 19.0 4 2.0
I want to add a column Feature_1 which, for each row of the dataframe, compute the median of column A for ALL the values which have the same Time value. This can be done as follows:
df['Feature_1'] = df.groupby('Time')['A'].transform(np.median)
print(df)
Output:
A Id Time Feature_1
0 10.0 1 0.0 10.0
1 15.0 1 1.0 15.0
2 NaN 1 2.0 19.0
3 11.0 2 0.0 10.0
4 16.0 2 1.0 15.0
5 25.0 2 2.0 19.0
6 10.0 3 0.0 10.0
7 15.0 3 2.0 19.0
8 9.0 4 0.0 10.0
9 14.0 4 1.0 15.0
10 19.0 4 2.0 19.0
My problem is now to compute another feature, Feature_2, which for each row of the dataframe, compute the median of column A for OTHER values which have the same Time value. I was not able to vectorize this, so my solution with a for loop:
df['feature_2'] = np.NaN
for i in range(len(df)):
current_Id = df.Id[i]
current_time = df.Time[i]
idx = (df.Time == current_time) & (df.Id != current_Id)
if idx.any():
df['feature_2'][i] = df.A[idx].median()
print(df)
Output:
A Id Time Feature_1 Feature_2
0 10.0 1 0.0 10.0 10.0
1 15.0 1 1.0 15.0 15.0
2 NaN 1 2.0 19.0 19.0
3 11.0 2 0.0 10.0 10.0
4 16.0 2 1.0 15.0 14.5
5 25.0 2 2.0 19.0 17.0
6 10.0 3 0.0 10.0 10.0
7 15.0 3 2.0 19.0 22.0
8 9.0 4 0.0 10.0 10.0
9 14.0 4 1.0 15.0 15.5
10 19.0 4 2.0 19.0 20.0
This is working but it is very slow as my dataframe has 1 million rows (but only four different IDs).
Is it possible to vectorize the creation of Feature_2 ?
I hope, I am clear enough. Live code can be found here.