pandas - 干净地创建新的分类数据列

Question

我可以像这样向 Pandas DataFrame 添加分类列：

import pandas as pd

label_type = pd.api.types.CategoricalDtype(categories=["positive", "negative"], ordered=False)

d = {'col1': [1, 2], 'col2': [3, 4]}
df = pd.DataFrame(data=d)

# Create a new column, setting the value universally to "positive"
df['label'] = pd.Series(["positive"] * len(df), dtype=label_type).values

这没有其他类型的速记那么优雅：

df['label2'] = "positive"  # sets entire column to str("positive")

但似乎底层类型只是str

print(type(df['label'].iloc[0]))
<class 'str'>

因此，熊猫似乎必须提前知道列类型。

有没有办法在不手动构建的情况下将分类列添加到数据框中Series？例如，

df['label3'] = label_type("positive")

Answer 1

这个怎么样：

df['col4'] = df.assign(col4 = 'positive')['col4'].astype(label_type)

df.info()

<class 'pandas.core.frame.DataFrame'>
RangeIndex: 2 entries, 0 to 1
Data columns (total 4 columns):
 #   Column  Non-Null Count  Dtype   
---  ------  --------------  -----   
 0   col1    2 non-null      int64   
 1   col2    2 non-null      int64   
 2   label   2 non-null      category
 3   col4    2 non-null      category
dtypes: category(2), int64(2)
memory usage: 412.0 bytes

虽然你仍然得到一个str type：

type(df['col4'].iloc[0])

str

因为我认为在这种情况下iloc[]将返回该类别的字符串表示形式。

或者只需分两步完成：

df['col4'] = 'positive'
df['col4'] = df['col4'].astype(label_type)