枢轴数据框与索引列中的列常量

Question

假设我有以下数据框，其中Y 和Z 在ID 中都是常量：

   ID TYPE  X    Y   Z
0   1    A  1  foo  10
1   1    B  2  foo  10
2   2    A  3  bar  20
3   2    B  4  bar  20
4   3    A  5  baz  30
5   3    B  6  baz  30

我想将数据从“长”格式重塑为“宽”格式：

   ID  XA  XB    Y   Z
0   1   1   2  foo  10
1   2   3   4  bar  20
2   3   5   6  baz  30

但是，如果我使用pandas.DataFrame.pivot():

df_new = df.pivot(index='ID', columns='TYPE')

我会得到Y 和Z 的重复：

      X       Y        Z    
TYPE  A  B    A    B   A   B
ID                          
1     1  2  foo  foo  10  10
2     3  4  bar  bar  20  20
3     5  6  baz  baz  30  30

要获得所需的输出，我可以执行以下操作：

import pandas as pd

df = pd.DataFrame({'ID': [1, 1, 2, 2, 3, 3],
                   'TYPE': ['A', 'B', 'A', 'B', 'A', 'B'],
                   'X': [1, 2, 3, 4, 5, 6],
                   'Y': ['foo', 'foo', 'bar', 'bar', 'baz', 'baz'],
                   'Z': [10, 10, 20, 20, 30, 30]})


def long_to_wide(df, i, j, varlist):
    df_wide = df.pivot(index='ID', columns='TYPE')
    df_wide.columns = [''.join(col).strip() for col in df_wide.columns.values]
    df_wide.reset_index(inplace=True)

    for var in varlist:
        if pd.Series.equals(df_wide[var + 'A'], df_wide[var + 'B']):
            df_wide.drop((var + 'B'), axis = 1, inplace = True)
        else:
            raise
            # Error handling of some sort...
        df_wide = df_wide.rename(columns={var + 'A': var})

    return df_wide


df_new = long_to_wide(df, 'ID', 'TYPE', ['Y', 'Z'])

但是，我觉得这一定是不必要的复杂。例如，要在 Stata 中获得所需的输出，可以运行：

reshape wide X, i(ID) j(TYPE)

或

reshape wide X, i(ID Y Z) j(TYPE)

这种情况很常见，因此我认为应该有一个内置的方法来处理它。但是在查看了Pandas 文档和 Stack Overflow 之后，我还没有找到更简单的解决方案。

有吗？

Answer 1

我刚刚对此进行了更好的了解，pandas.DataFrame.pivot() 函数实际上按预期执行。与Stata 的reshape 不同，后者是一个命令，在后台做了很多事情，pivot() 只是重新排列数据。

@Heleemur 的解决方案很聪明，效果很好，但通常重命名或删除重复项是您的责任。

这是一个基于pivot()（或pivot_table()）的直观解决方案：

import pandas as pd

df = pd.DataFrame({'ID': [1, 1, 2, 2, 3, 3],
                   'TYPE': ['A', 'B', 'A', 'B', 'A', 'B'],
                   'X': [1, 2, 3, 4, 5, 6],
                   'Y': ['foo', 'foo', 'bar', 'bar', 'baz', 'baz'],
                   'Z': [10, 10, 20, 20, 30, 30]})

wanted = df.pivot(index='ID', columns='TYPE')[[('X','A'), ('X','B'), ('Y','A'), ('Z','A')]].reset_index()
wanted.columns = wanted.columns.get_level_values(0)
wanted.columns = ['ID', 'XA', 'XB', 'Y', 'Z']
wanted

   ID  XA  XB    Y   Z
0   1   1   2  foo  10
1   2   3   4  bar  20
2   3   5   6  baz  30

另一种方式也如下：

wanted = df.pivot(index='ID', columns='TYPE').reset_index()
wanted.columns = [' '.join(col)for col in wanted.columns.values]
wanted = wanted.iloc[:, [0,2] + list(range(1, len(wanted.columns)-1, 2))]
wanted

   ID   X B  X A  Y A  Z A
0    1    2    1  foo   10
1    2    4    3  bar   20
2    3    6    5  baz   30

wanted.columns = ['ID', 'XB', 'XA', 'Y', 'Z']
wanted

   ID  XB  XA    Y   Z
0   1   2   1  foo  10
1   2   4   3  bar  20
2   3   6   5  baz  30

在具有更多列的较大数据框中，您可能希望保留原始名称。

---

编辑：

这是@Heleemur pivot_table() 提供的解决方案的等效解决方案：

wanted = df.pivot_table(index=['ID', 'Y', 'Z'], columns='TYPE').reset_index()
wanted.columns = [''.join(c) for c in wanted.columns.values]
wanted

   ID    Y   Z  XA  XB
0   1  foo  10   1   2
1   2  bar  20   3   4
2   3  baz  30   5   6

Answer 2

我通过设置索引、合并多索引列名和重置索引来做到这一点。我确信这也可以通过数据透视表实现（df 定义为您的示例数据框）。

df2 = df.set_index(['ID', 'Y', 'Z', 'TYPE']).unstack()
df2.columns = [''.join(c) for c in df2.columns.values]
df2.reset_index()

outputs:

   ID    Y   Z  XA  XB
0   1  foo  10   1   2
1   2  bar  20   3   4
2   3  baz  30   5   6