基于PySpark DataFrame生成行和为1的Transition Matrix

核心思路 #

通过频次统计→总行数计算→概率归一化→宽表转换四步，实现符合要求的转移矩阵：每个元素值为(from=X且to=Y的次数)/(from=X的总转移次数)，保证每行概率和为1。

具体实现代码 #

假设你的PySpark DataFrame名为df，执行以下操作：

1. 统计(from, to)对的出现频次

from pyspark.sql import functions as F

# 分组统计每个(from, to)组合的出现次数
count_df = df.groupBy("from", "to").agg(F.count("*").alias("count"))

2. 计算每个from的总转移次数并归一化概率
   用窗口函数计算每个from对应的总转移次数，再用频次除以总次数得到转移概率：

from pyspark.sql.window import Window

# 按from分组的窗口对象
from_window = Window.partitionBy("from")
# 计算总次数并生成归一化的转移概率
transition_df = count_df.withColumn("total_count", F.sum("count").over(from_window)) \
                        .withColumn("probability", F.col("count") / F.col("total_count"))

3. 转换为转移矩阵格式（宽表）
   通过pivot将to列的取值转为矩阵的列，缺失的转移路径填充为0：

# 转置为宽表结构，缺失值填充0
matrix_df = transition_df.groupBy("from") \
                         .pivot("to") \
                         .agg(F.first("probability", ignorenulls=True)) \
                         .fillna(0)

结果说明 #

执行完成后，matrix_df的每行对应一个from值，每列对应一个to值，单元格数值为转移概率，且每行概率和为1。基于你提供的输入数据，正确结果如下：

注：你提供的示例矩阵存在数据错误（比如from=2的转移只有to=4，概率应为1而非示例中的2/3），上述结果为基于输入数据的正确计算值。

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