熊猫融化后的分类列

Question

在pandas 中进行melt 操作后是否可以得到一个分类变量列？

如果我这样设置数据：

import pandas as pd
import numpy as np

df = pd.DataFrame(
    np.random.randn(3, 5), 
    columns=["A", "B", "C", "D", "E"]
)
df["id"] = range(1, 4)
df

|    |         A |         B |         C |         D |          E |   id |
|----|-----------|-----------|-----------|-----------|------------|------|
|  0 | -0.406174 | -0.686917 | -0.172913 | -0.273074 | -0.0246714 |    1 |
|  1 |  0.323783 | -1.7731   |  1.57581  | -1.15671  | -1.23926   |    2 |
|  2 | -1.1426   | -0.591279 |  1.15265  |  0.326712 | -0.86374   |    3 |

然后申请

melted_df = df.melt(id_vars="id", value_vars=["A", "B", "C", "D", "E"])
melted_df

|    |   id | variable   |      value |
|----|------|------------|------------|
|  0 |    1 | A          | -0.406174  |
|  1 |    2 | A          |  0.323783  |
|  2 |    3 | A          | -1.1426    |
|  3 |    1 | B          | -0.686917  |
|  4 |    2 | B          | -1.7731    |
|  5 |    3 | B          | -0.591279  |
|  6 |    1 | C          | -0.172913  |
|  7 |    2 | C          |  1.57581   |
|  8 |    3 | C          |  1.15265   |
|  9 |    1 | D          | -0.273074  |
| 10 |    2 | D          | -1.15671   |
| 11 |    3 | D          |  0.326712  |
| 12 |    1 | E          | -0.0246714 |
| 13 |    2 | E          | -1.23926   |
| 14 |    3 | E          | -0.86374   |

variable 列的 dtype 是 object

melted_df.dtypes

id            int64
variable     object
value       float64
dtype: object

我希望这是category。我知道，我可以通过以下方式轻松转换它：

melted_df["variable"].astype("category")

但对于大型数据集，我想避免这种开销。在documentation 中我没有找到这样的选项，但由于结果列包含定义的分类数据，我认为一定有可能。

Answer 1

我认为melt 不可能，因为当它创建该列时，它会推断出 dtype，而'category' 不是 pandas 当前推断的dtype。 （这是一个相关问题，它无法正确推断 Int32 dtypes Why is pandas.melt messing with my dtypes?）。

stack 如果您首先转换列，将保留分类 dtype。 stack 将导致与 melt 的排序略有不同，但数据将相同。 stack 在命名结果列时也有点笨拙。

df = df.set_index('id')
df.columns = df.columns.astype('category')

res = (df.stack()
         .rename_axis(['id', 'variable'])
         .rename('value')
         .reset_index())
#    id variable     value
#0    1        A  0.424781
#1    1        B -0.317107
#2    1        C  0.731121
#3    1        D  0.042642
#4    1        E  0.648352
#...
#13   3        D -0.889600
#14   3        E -1.822898

res.dtypes
#id             int64
#variable    category
#value        float64
#dtype: object