st19297 st19297 - 19 days ago 6
Python Question

Dropping rows in a multi-index data frame?

I have this df:

temp = pd.DataFrame({'tic': ['IBM', 'AAPL', 'AAPL', 'IBM', 'AAPL'],
'industry': ['A', 'B', 'B', 'A', 'B'],
'price': [np.nan, 5, 6, 11, np.nan],
'shares':[100, 60, np.nan, 100, np.nan],
'dates': pd.to_datetime(['1990-01-01', '1990-01-01', '1990-04-01',
'1990-04-01', '1990-08-01'])
})

temp.set_index(['tic', 'dates'], inplace=True)
temp.sort_index(inplace=True)


Which yields:

industry price shares
tic dates
AAPL 1990-01-01 B 5.0 60.0
1990-04-01 B 6.0 NaN
1990-08-01 B NaN NaN
IBM 1990-01-01 A NaN 100.0
1990-04-01 A 11.0 100.0


How can I create a
new column
in the data frame that shows the number of observations for each tic. So, the new column will like this:

New column
AAPL ... 3
... 3
... 3
IBM ... 2
... 2

Answer

you can use .groupby(level=0) and .filter() methods:

In [79]: temp.groupby(level=0).filter(lambda x: len(x) >= 3)
Out[79]:
                industry  price  shares
tic  dates
AAPL 1990-01-01        B    5.0    60.0
     1990-04-01        B    6.0     NaN
     1990-08-01        B    NaN     NaN

Answering your second question:

In [83]: temp['new'] = temp.groupby(level=0)['industry'].transform('size')

In [84]: temp
Out[84]:
                industry  price  shares  new
tic  dates
AAPL 1990-01-01        B    5.0    60.0    3
     1990-04-01        B    6.0     NaN    3
     1990-08-01        B    NaN     NaN    3
IBM  1990-01-01        A    NaN   100.0    2
     1990-04-01        A   11.0   100.0    2