Search code examples
pythonpandasindexingmaskargmax

`Pandas argmax` to get all `True` indices after masking (Python 3) (e.g. (pd.Series > 0).argmax()))

What I want to do but argmax only gives me the first value that is True:

Se = pd.Series(np.arange(6), index=list("abcdef"))
#a    0
#b    1
#c    2
#d    3
#e    4
#f    5
#dtype: int64

mask = (Se % 2 == 0)
#a     True
#b    False
#c     True
#d    False
#e     True
#f    False
#dtype: bool

mask.argmax()
#'a'

What I have to do instead:

Se[mask].index
# Index(['a', 'c', 'e'], dtype='object')

This isn't too inconvenient but I have to instantiate the Series first which reduces my productivity. It would be nice to be able to do this: 

(pd.Series(np.arange(6), index=list("abcdef")) % 2 == 0).argmax()

My question is: How can I do this using argmax? If this can't be done w/ argmax, could I do this w/ a different function in pandas?

Solution

  • You can use compress:

    idx = pd.Series(np.arange(6), index=list("abcdef")).compress(lambda x: x % 2 == 0).index

    The resulting output:

    Index(['a', 'c', 'e'], dtype='object')