Изменение формы датафреймов с мультииндексом

Я неоднократно пользовался методом reshape() библи NumPy, для быстрого создания датафреймов, которые использую для тех или инных примеров. Это действительно, очень удобный способ быстро и изменить форму массива:

arr = np.arange(12)
arr
array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11])

Мы просто указываем размер соответствующей оси и сразу получаем результат:

arr.reshape(6, 2)
array([[ 0,  1],
       [ 2,  3],
       [ 4,  5],
       [ 6,  7],
       [ 8,  9],
       [10, 11]])
arr.reshape(4, 3)
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11]])

А если в качестве размера одной из осей указать -1, то ее размер будет вычислен (если это возможно), на основании размера другой оси:

arr.reshape(2, -1)
array([[ 0,  1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10, 11]])

В общем удобно и крайне просто. Но снова возникает вопрос: как же быть с датафреймами, ведь помимо целочисленных, упорядоченных индексов, каждая ось (строки и столбцы) может быть помечена? Можем ли мы, например, так же легко поменять форму приведенного ниже датафрейма?

df = pd.DataFrame(np.arange(9).reshape(3, 3),
                  index=list('ABC'),
                  columns=list('XYZ'))
df
X Y Z
A 0 1 2
B 3 4 5
C 6 7 8

Очевидно, что поменять форму этого датафрейма так же просто, как мы недавно меняли форму массива NumPy у нас не получится. Так что же тогда понимается под сменой формы датафреймов? Ведь в каждом столбце мы храним признаки, а в каждой строке наблюдения! Т.е. это не просто массив чисел, это осмысленная таблица с данными.

На самом деле есть всего два варианта изменения формы датафреймов:

  • сделать столбцы строками;
  • сделать строки столбцами.

В первом случае говорят о "повороте", а взглянуть на результат такого поворота можно с помощью метода stack():

df1 = df.stack()
df1
A  X    0
   Y    1
   Z    2
B  X    3
   Y    4
   Z    5
C  X    6
   Y    7
   Z    8
dtype: int32

В результате мы получили иерархически проиндексированную серию (с мультииндексом).

Разумеется, если у нас есть серия с мультииндексом, то мы можем преобразовать строки в столбцы, т.е. выполнить обратный "поворот" и восстановить исходный датафрейм:

df1.unstack()
X Y Z
A 0 1 2
B 3 4 5
C 6 7 8

Все преобразования методами stack() и unstack() по умолчанию выполняются по самому внутреннему уровню иерархической индексации, но это поведение можно изменить если указать номер уровня:

df1.unstack(0)
A B C
X 0 3 6
Y 1 4 7
Z 2 5 8
df1.unstack(1)
X Y Z
A 0 1 2
B 3 4 5
C 6 7 8

Давайте приведем какой-нибудь более осмысленный пример, который помог бы проследить за всеми преобразованиями. Допустим у нас есть следующий датафрейм:

df = pd.DataFrame(np.random.randint(45, 60, (4, 4)).reshape(4, 4),
                 index=[['пловцы', 'пловцы', 'бегуны', 'бегуны'],
                         ['мальчики', 'девочки','мальчики', 'девочки']],
                 columns=[['группа 1', 'группа 1', 'группа 2', 'группа 2'],
                          ['рост', 'вес', 'рост', 'вес']])
df
группа 1 группа 2
рост вес рост вес
пловцы мальчики 52 51 49 47
девочки 58 46 55 55
бегуны мальчики 52 58 56 54
девочки 55 49 49 57

Давайте сначала взглянем на то как можно "повернуть" этот датафрейм по разным уровням мультииндекса. Сначала посмотрим, что будет если мы выполним поворот по самому внутреннему индексу:

df.stack()
группа 1 группа 2
пловцы мальчики вес 51 47
рост 52 49
девочки вес 46 55
рост 58 55
бегуны мальчики вес 58 54
рост 52 56
девочки вес 49 57
рост 55 49

Если мы укажем df.stack(1) то увидим точно такой же результат. Но что если мы хотим указать в каком порядке уровней столбцы должны укладываться в строки? В этом случае можно указать уровни в виде списка:

df.stack([0, 1])
пловцы  мальчики  группа 1  вес     51
                            рост    52
                  группа 2  вес     47
                            рост    49
        девочки   группа 1  вес     46
                            рост    58
                  группа 2  вес     55
                            рост    55
бегуны  мальчики  группа 1  вес     58
                            рост    52
                  группа 2  вес     54
                            рост    56
        девочки   группа 1  вес     49
                            рост    55
                  группа 2  вес     57
                            рост    49
dtype: int32
df.stack([1, 0])
пловцы  мальчики  вес   группа 1    51
                        группа 2    47
                  рост  группа 1    52
                        группа 2    49
        девочки   вес   группа 1    46
                        группа 2    55
                  рост  группа 1    58
                        группа 2    55
бегуны  мальчики  вес   группа 1    58
                        группа 2    54
                  рост  группа 1    52
                        группа 2    56
        девочки   вес   группа 1    49
                        группа 2    57
                  рост  группа 1    55
                        группа 2    49
dtype: int32

А теперь давайте рассмотрим, как ведет себя метод unstack():

df.unstack(0)
группа 1 группа 2
рост вес рост вес
бегуны пловцы бегуны пловцы бегуны пловцы бегуны пловцы
девочки 55 58 49 46 49 55 57 55
мальчики 52 52 58 51 56 49 54 47
df.unstack(1)
группа 1 группа 2
рост вес рост вес
девочки мальчики девочки мальчики девочки мальчики девочки мальчики
бегуны 55 52 49 58 49 56 57 54
пловцы 58 52 46 51 55 49 55 47
df.unstack([0, 1])
группа 1  рост  бегуны  девочки     55
                        мальчики    52
                пловцы  девочки     58
                        мальчики    52
          вес   бегуны  девочки     49
                        мальчики    58
                пловцы  девочки     46
                        мальчики    51
группа 2  рост  бегуны  девочки     49
                        мальчики    56
                пловцы  девочки     55
                        мальчики    49
          вес   бегуны  девочки     57
                        мальчики    54
                пловцы  девочки     55
                        мальчики    47
dtype: int32
df.unstack([1, 0])
группа 1  рост  девочки   бегуны    55
                          пловцы    58
                мальчики  бегуны    52
                          пловцы    52
          вес   девочки   бегуны    49
                          пловцы    46
                мальчики  бегуны    58
                          пловцы    51
группа 2  рост  девочки   бегуны    49
                          пловцы    55
                мальчики  бегуны    56
                          пловцы    49
          вес   девочки   бегуны    57
                          пловцы    55
                мальчики  бегуны    54
                          пловцы    47
dtype: int32

Методы stack() и unstack() могут оказаться очень полезны благодаря своей гибкости. Напоследок нужно упомянуть, что в некоторых ситуациях при обратном повороте в данных могут появиться NaN-ы. Допустим у нас есть вот такая серия:

ser = pd.Series(np.arange(7),
                index= [list('XXXXXYY'), list('ABCDEAB')])
ser
X  A    0
   B    1
   C    2
   D    3
   E    4
Y  A    5
   B    6
dtype: int32

Как видите на внешнем уровне в каждой подгруппе присутствует разное количество элементов, из за чего после обратного поворота, в результирующем датафрейме вместо недостающих элементов появятся NaN-ы:

ser.unstack()
A B C D E
X 0.0 1.0 2.0 3.0 4.0
Y 5.0 6.0 NaN NaN NaN

Тем не менее, эти NaN-ы вовсе не мешают восстановлению исходной серии:

ser.unstack().stack()
X  A    0.0
   B    1.0
   C    2.0
   D    3.0
   E    4.0
Y  A    5.0
   B    6.0
dtype: float64

Но если вам все же понадобилось сохранить NaN-ы при повороте датафрейма, то достаточно указать dropna=False

ser.unstack().stack(dropna=False)
X  A    0.0
   B    1.0
   C    2.0
   D    3.0
   E    4.0
Y  A    5.0
   B    6.0
   C    NaN
   D    NaN
   E    NaN
dtype: float64