pandas optimal way to do element-wise comparison between two series

user_3284713 0
Viewed 552

I have two pandas series:

s1 with potentially a very large number of rows and some NaN and s2 (which is a column in a dataframe (df) with only 20 rows.
The index of the two series differ.

s1:

id
1      4.5
2     15.0
3     13.0
4     14.0
5     18.0
6     15.0
7     13.0
8     14.0
9      NaN
10     NaN
11     NaN
12    18.0
13     NaN
14     NaN
15     NaN


df:

      col1    s2   
0     20.0    0.0
1     19.0    4.5
2     18.0    5.0
3     17.0    6.0
4     16.0    7.0
5     15.0    8.0
6     14.0    9.0
7     13.0   10.0
8     12.0   11.0
9     11.0   12.0
10    10.0   13.0
11     9.0   15.0
12     8.0   16.0
13     7.0   18.0
14     6.0   20.0
15     5.0   22.0
16     4.0   24.0
17     3.0   26.0
18     2.0   28.0
19     1.0  100.0

For each id of s1 I want to retrieve the value of col1 for the first element in s2 smaller or equal to id.

i.e. for id 1 we have s1 = 4.5, which is smaller or equal than df.s2 = 4.5, hence I want to retrieve the value 19. Accordingly, for id=2 in s1 I need to retrieve the value 9 in df.col1

This is my current solution. I am wondering if there is a better (faster, maybe a pandas function?) way to get the same result:

      output =  [min(df[df['s2'].le(element)].col1, default = np.NaN) for element in s1]

[19.0,
 9.0,
 10.0,
 10.0,
 7.0,
 9.0,
 10.0,
 10.0,
 nan,
 nan,
 nan,
 7.0,
 nan,
 nan,
 nan]
pandas dataframe python
2 Answers

Idea is use numpy and compare each value from column by each value of Series for 2d array, then pass to numpy.where, set NaN if no match and last use numpy.nanmean:

m = df['s2'].to_numpy() <= s1.to_numpy()[:, None]

a = np.nanmin(np.where(m, df['col1'], np.nan), axis=1)
print (a)
[19.  9. 10. 10.  7.  9. 10. 10. nan nan nan  7. nan nan nan]

Performance: original sample

In [63]: %%timeit
    ...: [min(df[df['s2'].le(element)].col1, default = np.NaN) for element in s1]
    ...: 
    ...: 
9.21 ms ± 305 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [64]: %%timeit
    ...: m = df['s2'].to_numpy() <= s1.to_numpy()[:, None]
    ...: a = np.nanmin(np.where(m, df['col1'], np.nan), axis=1)
72.4 µs ± 870 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Larger data 100 times:

#2k rows
df = pd.concat([df] * 100, ignore_index=True)
#1.5k rows
s1 = pd.concat([s1] * 100, ignore_index=True)


In [68]: %%timeit
    ...: [min(df[df['s2'].le(element)].col1, default = np.NaN) for element in s1]
    ...: 
    ...: 
1.12 s ± 17.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [69]: %%timeit
    ...: m = df['s2'].to_numpy() <= s1.to_numpy()[:, None]
    ...: a = np.nanmin(np.where(m, df['col1'], np.nan), axis=1)
34.2 ms ± 305 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
user_29010020

You can use interval index.

First the data: 

df1 = pd.DataFrame(
    np.array(
        [
            4.5,
            15.0,
            13.0,
            14.0,
            18.0,
            15.0,
            13.0,
            14.0,
            np.nan,
            np.nan,
            np.nan,
            18.0,
            np.nan,
            np.nan,
            np.nan,
        ]
    ),
    columns=["s1"],
)
print(df1)
       s1
0   4.500
1  15.000
2  13.000
3  14.000
4  18.000
5  15.000
6  13.000
7  14.000
8     nan
9     nan
10    nan
11 18.000
12    nan
13    nan
14    nan

Then the lookup dataframe: 

df = pd.DataFrame.from_dict(
    {
        "col1": {
            0: 20.0,
            1: 19.0,
            2: 18.0,
            3: 17.0,
            4: 16.0,
            5: 15.0,
            6: 14.0,
            7: 13.0,
            8: 12.0,
            9: 11.0,
            10: 10.0,
            11: 9.0,
            12: 8.0,
            13: 7.0,
            14: 6.0,
            15: 5.0,
            16: 4.0,
            17: 3.0,
            18: 2.0,
            19: 1.0,
        },
        "end": {
            0: 0.0,
            1: 4.5,
            2: 5.0,
            3: 6.0,
            4: 7.0,
            5: 8.0,
            6: 9.0,
            7: 10.0,
            8: 11.0,
            9: 12.0,
            10: 13.0,
            11: 15.0,
            12: 16.0,
            13: 18.0,
            14: 20.0,
            15: 22.0,
            16: 24.0,
            17: 26.0,
            18: 28.0,
            19: 100.0,
        },
    }
)
print(df)
    col1     end
0  20.000   0.000
1  19.000   4.500
2  18.000   5.000
3  17.000   6.000
4  16.000   7.000
5  15.000   8.000
6  14.000   9.000
7  13.000  10.000
8  12.000  11.000
9  11.000  12.000
10 10.000  13.000
11  9.000  15.000
12  8.000  16.000
13  7.000  18.000
14  6.000  20.000
15  5.000  22.000
16  4.000  24.000
17  3.000  26.000
18  2.000  28.000
19  1.000 100.000

Make a start column to create the interval, fill the first row with zero.

df["start"] = df["end"].shift().fillna(0)
print(df.head())
    col1   end  start
0 20.000 0.000  0.000
1 19.000 4.500  0.000
2 18.000 5.000  4.500
3 17.000 6.000  5.000
4 16.000 7.000  6.000

Create an interval index and set as index. 

idx = pd.IntervalIndex.from_arrays(df["start"], df["end"], closed="right")
df.index = idx
print(df.head())
             col1   end  start
(0.0, 0.0] 20.000 0.000  0.000
(0.0, 4.5] 19.000 4.500  0.000
(4.5, 5.0] 18.000 5.000  4.500
(5.0, 6.0] 17.000 6.000  5.000
(6.0, 7.0] 16.000 7.000  6.000

Final result

df1.loc[df1.dropna().index, "col1"] = df.loc[df1.loc[:, "s1"].dropna(), "col1"].values

print(df1)
      s1   col1
0   4.500 19.000
1  15.000  9.000
2  13.000 10.000
3  14.000  9.000
4  18.000  7.000
5  15.000  9.000
6  13.000 10.000
7  14.000  9.000
8     nan    nan
9     nan    nan
10    nan    nan
11 18.000  7.000
12    nan    nan
13    nan    nan
14    nan    nan

Full code without printouts. 

df["start"] = df["end"].shift().fillna(0)

idx = pd.IntervalIndex.from_arrays(df["start"], df["end"], closed="right")
df.index = idx

df1.loc[df1.dropna().index, "col1"] = df.loc[df1.loc[:, "s1"].dropna(), "col1"].values
user_108886550
Related