Select columns based on last value in pandas dataframe-CodePudding

I have a dataframe below as an example. I want to iterate through each wag number and select only columns A or B based on its last value. For example if we look at wag == 1 then last value of A == 3 whereas last value of B == 34 so then I only interested in column A as its value is less and so on through all wags. Finally, I want to have a list as below:

foo_list = [[10, 8, 7, 3], [15, 12, 5], [1, 2, 2, 1]]

import pandas as pd

foo = pd.DataFrame({
                    'wag' : [1, 1, 1, 1, 44, 44, 44, 56, 56, 56, 56],
                     'A' :  [10, 8, 7, 3, 13, 22, 18, 43, 40, 38, 12],
                     'B' : [55, 48, 37, 34, 15, 12, 5, 1, 2, 2, 1]
                     })

In reality, it can be many column to compare not only 2 as I showed. And all the time I need to store in list columns with minimum last value amoung others. I thought about for loop, but cannot think how to compare columns with each other and then select the one I want.

CodePudding user response：

For each group get minimal last values and convert this column to list, here is set wag to index first for processing all another columns like A,B:

f = lambda x: list(x.loc[:, x.iloc[-1].idxmin()])
L = foo.set_index('wag').groupby('wag', sort=False).apply(f).tolist()
print (L)
[[10, 8, 7, 3], [15, 12, 5], [1, 2, 2, 1]]