示例

import pandas as pd
import numpy as np
from sklearn.preprocessing import OneHotEncoder


def onehot_encoding(X_train, X_test, cat_cols, max_categories=20, sparse_threshold=100000):
    """
    OneHot编码（增强版）

    Parameters:
    -----------
    X_train : DataFrame
        训练集特征数据
    X_test : DataFrame
        测试集特征数据
    cat_cols : list
        需要进行编码的分类特征列名列表
    max_categories : int, default=20
        最大类别数阈值，超过此值的特征将不被编码
    sparse_threshold : int, default=100000
        稀疏矩阵阈值，当样本数超过此值时使用稀疏输出

    Returns:
    --------
    X_train : DataFrame
        编码后的训练集
    X_test : DataFrame
        编码后的测试集
    encoder : OneHotEncoder
        训练好的编码器对象
    encoded_cols : list
        实际被编码的特征列名
    high_card_cols : list
        因类别过多而未编码的特征列名
    """

    # 1. 筛选高基数特征并给出警告
    high_card_cols = []
    cols_to_encode = []

    for col in cat_cols:
        n_unique = X_train[col].nunique()
        if n_unique <= max_categories:
            cols_to_encode.append(col)
        else:
            high_card_cols.append(col)
            print(f"⚠️ 警告：特征 '{col}' 有 {n_unique} 个类别（超过阈值 {max_categories}），未被编码")

    # 2. 如果没有需要编码的特征，直接返回
    if not cols_to_encode:
        print("ℹ️ 没有符合条件的特征需要编码")
        return X_train, X_test, None, [], high_card_cols

    # 3. 根据数据量决定是否使用稀疏矩阵
    n_samples = len(X_train)
    sparse_output = n_samples > sparse_threshold
    if sparse_output:
        print(f"ℹ️ 样本数 {n_samples} 超过阈值 {sparse_threshold}，使用稀疏矩阵输出")
    else:
        print(f"ℹ️ 样本数 {n_samples} 未超过阈值 {sparse_threshold}，使用密集矩阵输出")

    # 4. 初始化编码器
    encoder = OneHotEncoder(
        handle_unknown='ignore',
        sparse_output=sparse_output
    )

    # 5. 执行编码
    encoded_train = encoder.fit_transform(X_train[cols_to_encode])
    encoded_test = encoder.transform(X_test[cols_to_encode])

    # 6. 生成特征名称并创建DataFrame
    feature_names = encoder.get_feature_names_out(cols_to_encode)

    if sparse_output:
        # 稀疏矩阵需要转换为密集格式才能创建DataFrame
        # 注意：这里为了兼容性，如果数据量特别大，建议直接使用稀疏格式
        encoded_train_dense = encoded_train.toarray() if hasattr(encoded_train, 'toarray') else encoded_train
        encoded_test_dense = encoded_test.toarray() if hasattr(encoded_test, 'toarray') else encoded_test

        df_encoded_train = pd.DataFrame(
            encoded_train_dense,
            columns=feature_names,
            index=X_train.index
        )
        df_encoded_test = pd.DataFrame(
            encoded_test_dense,
            columns=feature_names,
            index=X_test.index
        )
    else:
        df_encoded_train = pd.DataFrame(
            encoded_train,
            columns=feature_names,
            index=X_train.index
        )
        df_encoded_test = pd.DataFrame(
            encoded_test,
            columns=feature_names,
            index=X_test.index
        )

    # 7. 替换原始特征
    X_train = X_train.drop(cols_to_encode, axis=1).join(df_encoded_train)
    X_test = X_test.drop(cols_to_encode, axis=1).join(df_encoded_test)

    # 8. 输出编码统计信息
    print(f"✅ 成功编码 {len(cols_to_encode)} 个特征，生成 {len(feature_names)} 个新特征")
    print(f"📊 编码特征: {cols_to_encode}")
    if high_card_cols:
        print(f"⏭️ 未编码特征: {high_card_cols}")

    return X_train, X_test, encoder, cols_to_encode, high_card_cols

调用示例

# ============ 使用示例 ============
if __name__ == "__main__":
    # 创建示例数据
    train = pd.DataFrame({
        '颜色': ['红', '蓝', '红', '绿', '黄', '红'],
        '城市': ['北京', '上海', '北京', '广州', '深圳', '杭州'],
        '价格': [100, 200, 150, 180, 220, 190],
        '用户ID': ['U001', 'U002', 'U003', 'U004', 'U005', 'U006']
    })

    test = pd.DataFrame({
        '颜色': ['蓝', '黄', '红', '紫'],  # '紫'是训练集未见的类别
        '城市': ['上海', '深圳', '北京', '成都'],
        '价格': [190, 210, 160, 175],
        '用户ID': ['U007', 'U008', 'U009', 'U010']
    })

    # 执行编码（设置较小的阈值以演示警告）
    train_enc, test_enc, encoder, encoded_cols, high_card_cols = onehot_encoding(
        X_train=train,
        X_test=test,
        cat_cols=['颜色', '城市', '用户ID'],
        max_categories=3,  # 故意设小以演示警告
        sparse_threshold=100  # 故意设小以演示稀疏矩阵
    )

    print("\n" + "=" * 50)
    print("编码后的训练集:")
    print(train_enc)
    print("\n编码后的测试集:")
    print(test_enc)

    # 验证未见类别的处理
    print("\n" + "=" * 50)
    print("验证未见类别处理（测试集中的'紫'和'成都'）：")
    print("注意：'颜色_紫'和'城市_成都'全部为0，表示OneHot编码正确处理了未知类别")

稀疏矩阵原理详解

1. 什么是稀疏矩阵？

稀疏矩阵是指矩阵中大部分元素为0的矩阵。在OneHot编码中，每个样本只有一个类别值为1，其余都是0，天然适合用稀疏格式存储。

python

# 假设有100个城市，100万条数据
# 密集矩阵：100万 × 100 = 1亿个元素
# 稀疏矩阵：只存储非零元素的位置和值

2. 稀疏矩阵的存储格式

SciPy提供了多种稀疏矩阵格式，各有优劣：

(1) CSR格式（Compressed Sparse Row）- 最常用

from scipy.sparse import csr_matrix
import numpy as np

# 原始密集矩阵
dense = np.array([
    [1, 0, 0, 0],
    [0, 2, 0, 0],
    [0, 0, 3, 0],
    [0, 0, 0, 4]
])

sparse = csr_matrix(dense)

# CSR存储三个数组：
# data: 存储所有非零值 → [1, 2, 3, 4]
# indices: 存储每个非零值的列索引 → [0, 1, 2, 3]
# indptr: 存储每行起始位置 → [0, 1, 2, 3, 4]

CSR格式示意图：

原始矩阵:
[1, 0, 0, 0]  ← row 0
[0, 2, 0, 0]  ← row 1
[0, 0, 3, 0]  ← row 2
[0, 0, 0, 4]  ← row 3

CSR存储:
data    = [1, 2, 3, 4]        # 所有非零值
indices = [0, 1, 2, 3]        # 每个值所在的列
indptr  = [0, 1, 2, 3, 4]     # 每行的起始位置

(2) COO格式（Coordinate Format）- 直观但内存较大

from scipy.sparse import coo_matrix

# COO存储三个数组：
# data: [1, 2, 3, 4]
# row: [0, 1, 2, 3]    # 行索引
# col: [0, 1, 2, 3]    # 列索引

(3) CSC格式（Compressed Sparse Column）- 列访问高效

from scipy.sparse import csc_matrix
# 类似于CSR，但按列压缩

3. OneHot编码中的稀疏矩阵

from sklearn.preprocessing import OneHotEncoder
import pandas as pd
from scipy.sparse import csr_matrix

# 创建数据
data = pd.DataFrame({
    '颜色': ['红', '蓝', '红', '绿', '黄', '红', '蓝', '绿'],
    '城市': ['北京', '上海', '北京', '广州', '深圳', '杭州', '上海', '广州']
})

# 使用稀疏输出
encoder = OneHotEncoder(sparse_output=True)
encoded = encoder.fit_transform(data[['颜色', '城市']])

print(f"形状: {encoded.shape}")
print(f"非零元素数量: {encoded.nnz}")  # nnz = number of non-zero
print(f"总元素数量: {encoded.shape[0] * encoded.shape[1]}")
print(f"稀疏度: {1 - encoded.nnz / (encoded.shape[0] * encoded.shape[1]):.2%}")