如何用内置工具将Pandas DataFrame转换为指定稀疏矩阵？

给定如下Pandas DataFrame：

data = {'col_1': ['a', 'b'], 'col_2': ['b', 'c']}
df = pd.DataFrame(data)

需要将其转换为如下的numpy稀疏表示形式：

[[[1, 0], [0, 0]], [[0, 1], [1, 0]], [[0, 0], [0, 1]]]

其中每个2x2矩阵分别对应原DataFrame中'a'、'b'、'c'的出现情况。当前已有繁琐的实现方式，希望找到Pandas、Polars、Numpy、TensorFlow的内置/标准方法完成转换。

一、Numpy 高效实现 #

利用Numpy向量化操作直接生成目标张量，是效率最高的方案之一：

import numpy as np
import pandas as pd

data = {'col_1': ['a', 'b'], 'col_2': ['b', 'c']}
df = pd.DataFrame(data)

# 获取所有唯一标签并建立索引映射
labels = np.unique(df.values)
n_labels = len(labels)
label_to_idx = {v:i for i,v in enumerate(labels)}

# 将DataFrame的字符值转换为索引矩阵
idx_matrix = np.vectorize(label_to_idx.get)(df.values)

# 生成目标稀疏张量，调整维度匹配需求
result = np.eye(n_labels)[idx_matrix.T].transpose(1,0,2)

print(result.tolist())

二、Pandas 优化实现 #

结合pd.Categorical统一类别，再配合Numpy操作简化流程：

import pandas as pd
import numpy as np

data = {'col_1': ['a', 'b'], 'col_2': ['b', 'c']}
df = pd.DataFrame(data)

# 统一所有列的分类类别
all_labels = np.unique(df.values)
cat_df = df.apply(lambda x: pd.Categorical(x, categories=all_labels))

# 获取类别索引矩阵并生成目标张量
idx_matrix = cat_df.values.codes
n_cats = len(all_labels)
result = np.eye(n_cats)[idx_matrix.T].transpose(1,0,2).tolist()

print(result)

三、Polars 实现 #

借助Polars的分类类型和向量化操作完成转换：

import polars as pl
import numpy as np

data = {'col_1': ['a', 'b'], 'col_2': ['b', 'c']}
df = pl.DataFrame(data)

# 获取全局唯一标签
labels = np.unique(df.to_numpy().flatten())
n_labels = len(labels)

# 将DataFrame转换为类别索引矩阵
idx_matrix = df.select(pl.all().cast(pl.Categorical(categories=labels))).to_numpy().astype(int)

# 生成目标稀疏张量
result = np.eye(n_labels)[idx_matrix.T].transpose(1,0,2).tolist()

print(result)

四、TensorFlow 实现 #

利用TensorFlow的tf.one_hot函数完成转换，适合深度学习场景：

import tensorflow as tf
import pandas as pd

data = {'col_1': ['a', 'b'], 'col_2': ['b', 'c']}
df = pd.DataFrame(data)

# 建立标签到索引的映射
labels = pd.unique(df.values.ravel())
label_to_idx = {v:i for i,v in enumerate(labels)}
n_labels = len(labels)

# 转换为TensorFlow索引张量
idx_tensor = tf.convert_to_tensor(df.replace(label_to_idx).values, dtype=tf.int32)

# 生成one-hot张量并调整维度
one_hot = tf.one_hot(idx_tensor, depth=n_labels)
result = tf.transpose(one_hot, perm=[2,0,1]).numpy().tolist()

print(result)

内容的提问来源于stack exchange，提问作者user17033672