pandas DataFrame: get column item when the corresponding item in another column is greater than a value

Question

I have the following pandas dataFrame. It is large dataframe with more than 500k rows.

    Event_Number  Well  p_and_s
0              1     7      4.0
1              1     9      0.0
2              1    15      0.0
3              2     7      2.0
4              2     9      7.0
5              2    15      0.0
6              3     5      0.0
7              3     7      8.0
8              3    16      3.0
9              4     7      8.0
10             4    15      0.0
11             5     7      8.0
12             5     9      3.0
13             5    15      6.0
14             6     5      0.0
15             6     7      8.0
16             7     7      8.0
17             7     9      0.0
18             7    15      0.0
19             8     7      8.0
20             8    15      4.0

I want to find for each group of [column: Event_Number] what [column: Well] has greater than 2 value in column [p_and_s].

The final dataFrame should look like this with a new column listing all well having p_and_s greater than 2

        Event_Number  Well  p_and_s  well_array
0                 1     7      4.0  [7]
1                 1     9      0.0  [7]
2                 1    15      0.0  [7]
3                 2     7      2.0  [9]
4                 2     9      7.0  [9]
5                 2    15      0.0  [9]
6                 3     5      0.0  [7, 16]
7                 3     7      8.0  [7, 16]
8                 3    16      3.0  [7, 16]
9                 4     7      8.0  [7]
10                4    15      0.0  [7]
11                5     7      8.0  [7, 9, 15]
12                5     9      3.0  [7, 9, 15]
13                5    15      6.0  [7, 9, 15]
14                6     5      0.0  [7]
15                6     7      8.0  [7]
16                7     7      8.0  [7]
17                7     9      0.0  [7]
18                7    15      0.0  [7]
19                8     7      8.0  [7, 15]
20                8    15      4.0  [7, 15]

Answer 1

Here is one way.

s = df[df['p_and_s'] > 2].groupby('Event_Number')['Well'].apply(list)
df['well_array'] = df['Event_Number'].map(s)

Explanation

• Create a series mapping Event_Number to Well after applying a filter on p_and_s.
• Map to original dataframe via pd.Series.map.
• For performance, where possible lambda functions should be avoided as they represent expensive implicit loops.

Result

    Event_Number  Well  p_and_s  well_array
0              1     7      4.0         [7]
1              1     9      0.0         [7]
2              1    15      0.0         [7]
3              2     7      2.0         [9]
4              2     9      7.0         [9]
5              2    15      0.0         [9]
6              3     5      0.0     [7, 16]
7              3     7      8.0     [7, 16]
8              3    16      3.0     [7, 16]
9              4     7      8.0         [7]
10             4    15      0.0         [7]
11             5     7      8.0  [7, 9, 15]
12             5     9      3.0  [7, 9, 15]
13             5    15      6.0  [7, 9, 15]
14             6     5      0.0         [7]
15             6     7      8.0         [7]
16             7     7      8.0         [7]
17             7     9      0.0         [7]
18             7    15      0.0         [7]
19             8     7      8.0     [7, 15]
20             8    15      4.0     [7, 15]

Answer 2

You can try this:

df.groupby('Event_Number')\
  .apply(lambda x: x.loc[x.p_and_s > 2,'Well'].tolist())\
  .to_frame(name='well_array')\
  .merge(df, right_on=['Event_Number'], left_index=True)

Output:

    well_array  Event_Number  Well  p_and_s
0          [7]             1     7      4.0
1          [7]             1     9      0.0
2          [7]             1    15      0.0
3          [9]             2     7      2.0
4          [9]             2     9      7.0
5          [9]             2    15      0.0
6      [7, 16]             3     5      0.0
7      [7, 16]             3     7      8.0
8      [7, 16]             3    16      3.0
9          [7]             4     7      8.0
10         [7]             4    15      0.0
11  [7, 9, 15]             5     7      8.0
12  [7, 9, 15]             5     9      3.0
13  [7, 9, 15]             5    15      6.0
14         [7]             6     5      0.0
15         [7]             6     7      8.0
16         [7]             7     7      8.0
17         [7]             7     9      0.0
18         [7]             7    15      0.0
19     [7, 15]             8     7      8.0
20     [7, 15]             8    15      4.0