pandas DataFrame .stack (dropna = False), но с сохранением существующих комбинаций уровней

Question

Мои данные выглядят так:

import numpy as np
import pandas as pd

# My Data
enroll_year = np.arange(2010, 2015)
grad_year = enroll_year + 4
n_students = [[100, 100, 110, 110, np.nan]]

df = pd.DataFrame(
    n_students, 
    columns=pd.MultiIndex.from_arrays(
        [enroll_year, grad_year], 
        names=['enroll_year', 'grad_year']))

print(df)
# enroll_year 2010 2011 2012 2013 2014
# grad_year   2014 2015 2016 2017 2018
# 0            100  100  110  110  NaN

Я пытаюсь собрать данные, один уровень столбца / индекса для года регистрации, один для года выпуска и один для чиселстуденты, которые должны выглядеть как

# enroll_year  grad_year    n
# 2010         2014         100.0
# .            .                .
# .            .                .
# .            .                .
# 2014         2018           NaN

Данные, полученные с помощью .stack(), очень близки, но пропущенные записи пропущены,

df1 = df.stack(['enroll_year', 'grad_year'])
df1.index = df1.index.droplevel(0)
print(df1)
# enroll_year  grad_year
# 2010         2014         100.0
# 2011         2015         100.0
# 2012         2016         110.0
# 2013         2017         110.0
# dtype: float64

Итак, .stack(dropna=False)пробуется, но он расширит уровни индекса до всех комбинаций лет регистрации и окончания обучения

df2 = df.stack(['enroll_year', 'grad_year'], dropna=False)
df2.index = df2.index.droplevel(0)
print(df2)
# enroll_year  grad_year
# 2010         2014         100.0
#              2015           NaN
#              2016           NaN
#              2017           NaN
#              2018           NaN
# 2011         2014           NaN
#              2015         100.0
#              2016           NaN
#              2017           NaN
#              2018           NaN
# 2012         2014           NaN
#              2015           NaN
#              2016         110.0
#              2017           NaN
#              2018           NaN
# 2013         2014           NaN
#              2015           NaN
#              2016           NaN
#              2017         110.0
#              2018           NaN
# 2014         2014           NaN
#              2015           NaN
#              2016           NaN
#              2017           NaN
#              2018           NaN
# dtype: float64

И мне нужно установить подмножество df2, чтобы получить желаемый набор данных.

existing_combn = list(zip(
    df.columns.levels[0][df.columns.labels[0]], 
    df.columns.levels[1][df.columns.labels[1]]))

df3 = df2.loc[existing_combn]
print(df3)
# enroll_year  grad_year
# 2010         2014         100.0
# 2011         2015         100.0
# 2012         2016         110.0
# 2013         2017         110.0
# 2014         2018           NaN
# dtype: float64

Хотя это добавляет к моему коду только несколько дополнительных строк, мне интересно, есть ли какие-нибудь более подходящие и аккуратные подходы.

Answer 1

Используйте unstack с pd.DataFrame, затем reset_index и drop ненужные столбцы и rename столбец как:

pd.DataFrame(df.unstack()).reset_index().drop('level_2',axis=1).rename(columns={0:'n'})

   enroll_year  grad_year      n
0         2010       2014  100.0
1         2011       2015  100.0
2         2012       2016  110.0
3         2013       2017  110.0
4         2014       2018    NaN

Или:

df.unstack().reset_index(level=2, drop=True)
enroll_year  grad_year
2010         2014         100.0
2011         2015         100.0
2012         2016         110.0
2013         2017         110.0
2014         2018           NaN
dtype: float64

Или:

df.unstack().reset_index(level=2, drop=True).reset_index().rename(columns={0:'n'})
   enroll_year  grad_year      n
0         2010       2014  100.0
1         2011       2015  100.0
2         2012       2016  110.0
3         2013       2017  110.0
4         2014       2018    NaN

---

Объяснение:

print(pd.DataFrame(df.unstack()))
                             0
enroll_year grad_year         
2010        2014      0  100.0
2011        2015      0  100.0
2012        2016      0  110.0
2013        2017      0  110.0
2014        2018      0    NaN

print(pd.DataFrame(df.unstack()).reset_index().drop('level_2',axis=1))
   enroll_year  grad_year      0
0         2010       2014  100.0
1         2011       2015  100.0
2         2012       2016  110.0
3         2013       2017  110.0
4         2014       2018    NaN

print(pd.DataFrame(df.unstack()).reset_index().drop('level_2',axis=1).rename(columns={0:'n'}))
   enroll_year  grad_year      n
0         2010       2014  100.0
1         2011       2015  100.0
2         2012       2016  110.0
3         2013       2017  110.0
4         2014       2018    NaN