如何从索引DataFrame高效生成行对二元矩阵？

给定如下由元组构成的Pandas DataFrame：

import pandas as pd

df = pd.DataFrame({
    1: [(), (1, 2), (1, 2, 5, 7), (1, 2),    (1, 2, 5), (1, 2)],
    2: [(), (1, 2), (1, 2, 5, 7), (1, 2),    (1, 2, 5), (1, 2)],
    3: [(), (3, 4), (3, 4),       (3, 5, 6), (3, 4, 6), (3, 4, 5, 7)],
    4: [(), (3, 4), (3, 4),       (),        (3, 4, 6), (3, 4, 5, 7)],
    5: [(), (),     (1, 2, 5, 7), (3, 5, 6), (1, 2, 5), ()],
    6: [(), (),     (),           (3, 5, 6), (3, 4, 6), ()],
    7: [(), (),     (1, 2, 5, 7), (),        (),        (3, 4, 5, 7)],
    8: [(), (),     (),           (),        (),        ()],
    9: [(), (),     (),           (),        (),        ()]
})

需高效生成目标输出：

• 每行对应一对元素对（如示例中的(1,2)）
• 每列对应原DataFrame的列（1-9）
• 若该元素对同时存在于原列的任意元组子集中，则标记1，否则标记0。例如元素对(1,2)除第一列外在其余列的子集中均存在，对应值为[0,1,1,1,1,1,...]

需实现高效Python函数，并了解Pandas、TensorFlow、Spark是否有现成处理方案。

一、高效Python实现 #

核心逻辑为：先提取所有出现的元素并生成元素对，再将原DataFrame的每个元组转为集合（优化子集检测性能），最后遍历每个元素对和列，判断是否存在包含该元素对的子集。

import pandas as pd
from itertools import combinations

def generate_element_pair_matrix(df):
    # 提取所有唯一元素
    all_elements = set()
    for col in df.columns:
        for tpl in df[col]:
            all_elements.update(tpl)
    all_elements = sorted(all_elements)
    
    # 生成所有无序元素对
    pairs = list(combinations(all_elements, 2))
    
    # 预处理列数据：将元组转为集合，减少重复计算
    col_sets = {col: [set(tpl) for tpl in df[col]] for col in df.columns}
    
    # 生成结果矩阵
    result = []
    target_pair = set()
    for a, b in pairs:
        target_pair.clear()
        target_pair.update((a, b))
        row = []
        for col in df.columns:
            # 检查列中是否有子集包含目标元素对
            has_pair = any(target_pair.issubset(s) for s in col_sets[col])
            row.append(1 if has_pair else 0)
        result.append(row)
    
    # 转为DataFrame格式
    return pd.DataFrame(
        result,
        index=[f"{p[0]},{p[1]}" for p in pairs],
        columns=df.columns
    )

若需求为生成行索引对的检测矩阵，只需修改元素对的生成逻辑为行索引组合，并调整判断条件（如判断两行元组是否相等、是否有交集等），核心优化思路一致。

二、Pandas简化实现 #

Pandas无直接对应API，但可通过向量化操作与自定义函数简化代码，中等规模数据下效率足够：

import pandas as pd
from itertools import combinations

def pandas_element_pair_matrix(df):
    # 提取所有唯一元素
    all_elements = set(df.stack().explode().dropna().unique())
    pairs = list(combinations(all_elements, 2))
    
    # 定义列级检测函数
    def check_col(col, target):
        return col.apply(lambda x: target.issubset(set(x))).any().astype(int)
    
    # 生成结果
    result_rows = []
    for pair in pairs:
        target = set(pair)
        result_rows.append(df.apply(check_col, target=target).tolist())
    
    return pd.DataFrame(
        result_rows,
        index=[f"{p[0]},{p[1]}" for p in pairs],
        columns=df.columns
    )

若需处理大规模数据，可结合swifter库加速apply操作，它会自动选择向量化或并行执行方式优化性能。

三、TensorFlow/Spark方案 #

TensorFlow #

TensorFlow专注于深度学习任务，无针对此类集合检测的现成API。若强行实现，需将数据转为张量并编写自定义运算，但这类场景下TF的灵活性与效率均不如Pandas，不推荐使用。

Spark #

Spark适用于超大规模分布式数据处理，可通过两种方式实现：

1. RDD方式：将列数据转为RDD，预处理元组为集合，生成元素对后广播至所有节点，遍历列RDD检测子集包含关系，最后汇总结果。
2. DataFrame方式：定义UDF实现集合包含判断，通过Spark SQL的分组聚合完成检测。

Spark的优势在于处理单机无法承载的超大规模数据，但中小规模数据下，Pandas实现更简洁高效。

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