如何用Numpy计算多组任意尺寸列联表的期望值（无循环）

目标 #

通过行和乘以列和计算多个列联表的期望值Eij。

输入示例 #

输入为一组三维Numpy数组，示例如下：

import numpy as np
mat = np.array(
    [[[11., 13.],
     [12., 14.]],

     [[ 8., 10.],
      [15., 17.]],

     [[11., 10.],
     [12., 17.]]])

现有问题 #

现有代码仅能实现2x2矩阵的期望值计算，需扩展至任意尺寸的矩阵，且输出维度与输入完全一致；同时因涉及大规模数组，要求完全采用Numpy实现，避免使用循环。

现有代码 #

cols = np.sum(mat, axis=1)
rows = np.sum(mat, axis=2)
tots = np.sum(cols, 1)
exp_00 = cols[:,0]*rows[:,0]/tots
exp_01 = cols[:,1]*rows[:,0]/tots
exp_10 = cols[:,0]*rows[:,1]/tots
exp_11 = cols[:,1]*rows[:,1]/tots
mat_exp = np.array([exp_00, exp_01, exp_10, exp_11]).T.reshape(len(mat),2,2)

print(mat_exp)

现有输出 #

[[[11.04 12.96]
  [11.96 14.04]]

 [[ 8.28  9.72]
  [14.72 17.28]]

 [[ 9.66 11.34]
  [13.34 15.66]]]

解决方案 #

利用Numpy的广播和矩阵乘法特性，无需循环即可实现任意尺寸列联表的期望值计算，代码如下：

import numpy as np

def calculate_expected_values(mat):
    # 计算每个列联表的行和（保留维度以支持广播）
    row_sums = np.sum(mat, axis=2, keepdims=True)
    # 计算每个列联表的列和（保留维度以支持广播）
    col_sums = np.sum(mat, axis=1, keepdims=True)
    # 计算每个列联表的总样本量
    total_sums = np.sum(row_sums, axis=1, keepdims=True)
    # 通过矩阵乘法实现行和与列和的外积，再除以总样本量得到期望值
    expected = (row_sums @ col_sums) / total_sums
    return expected

# 测试示例
mat = np.array(
    [[[11., 13.],
     [12., 14.]],

     [[ 8., 10.],
      [15., 17.]],

     [[11., 10.],
     [12., 17.]]])

print(calculate_expected_values(mat))

代码说明 #

1. keepdims=True 保留求和后的维度，确保后续运算能匹配输入的三维结构（形状为(n_tables, n_rows, n_cols)）
2. row_sums @ col_sums 本质是对每个列联表的行和向量与列和向量做外积，生成和原矩阵同尺寸的中间矩阵
3. 最后除以总样本量，得到每个单元格的期望值Eij

运行结果与现有代码输出完全一致，且支持任意M×N尺寸的列联表。

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