CODEWITHSUNDEEP
pythonpandasdataframemulti-index
Amenhotep
Amenhotep

Reputation: 919

Get boolean array (or generator) from pandas DataFrame index

I have an indexed pandas DataFrame similar to this:

import pandas as pd

df = pd.DataFrame({'type':['good','good','bad'], 'nr':[0,1,2], 'value':[1,2,3]})

df.set_index(['type','nr'], inplace=True)

df
Out[153]:

        | value
type nr |
--------+-------
good  0 |   1
      1 |   2
--------+-------
bad   2 |   3

I want to plot the values (with matplotlib) together with an indicator 'good'/'bad'. For this indicator, a 0/1 numpy array would suffice -- 0 where the data is 'bad' and 1 where the data is 'good'. But the dataset is quite large and I would prefer a generator, not an actual array.

I don't know how to create this array/generator from df. I need to use it like this:

import matplotlib.pyplot as plt

plt.figure()
plt.plot(good_or_bad_indicator)
plt.plot(df)
plt.show()

Could someone please help?

PS: I would prefer a solution based on a lambda function, because in my real case there are more than two values in the 'type' column and I might want to design a more complex indicator to be plotted.

Upvotes: 0

Views: 797

Answers (1)

tktk
tktk

Reputation: 11744

There are a few methods to get your indicator:

df.index.get_level_values(0) == 'good'

is the simplest. Also check out isin if you have more than one "good" option. You can also iterate over the index, if you want something more complex (instead of x[0] == put a lambda or a method on x):

[x[0] == 'good' for x in df.index]

If you need the entire row, you can use df.iterrows()

Upvotes: 2

Related Questions