如何加速Python中多层嵌套for循环的矩阵函数？

我想了解是否存在其他方法来加速矩阵函数中的多层嵌套for循环。以下是我的函数：

def matrix(Xbin, y):
    labels = np.unique(y)
    con_matrix = []
    start = time.time()
    for i in range(len(labels)):
        for j in range(i + 1, len(labels)):
            # Crossover
            for u in Xbin[y == labels[i]]:
                for v in Xbin[y == labels[j]]:
                    con_matrix.append(np.bitwise_xor(u, v))
    end = time.time()
    duration = end - start
    print("total time for nested loop: ", duration)
    constraint_matrix = np.array(con_matrix)
    bin_attr_dim = [i for i in range(1, Xbin.shape[1] + 1)]

    df = pd.DataFrame(constraint_matrix, columns=bin_attr_dim)

    return df

背景说明 #

• Xbin 是numpy.ndarray类型，每行对应一个二值化样本，每列对应一个二值属性；y 是分组标签（例如1、2、3）。
• 函数功能：生成不同分组样本两两之间的按位异或结果，最终返回DataFrame格式的结果。
• 性能问题：在鸢尾花这类小数据集上运行速度很快，但在乳腺癌数据集上耗时明显增加，更大数据集下性能瓶颈会更突出。

生成二值化数据集的代码如下：

def binarize_dataset(X, y):
    cutpoints = {}

    att = -1
    for row in X.T:
        att += 1
        labels = None  # Previous labels
        u = -9999  # Previous xi

        # Finding transitions
        for v in sorted(np.unique(row)):
            variation = v - u  # Current - Previous

            # Classes where current v appears
            indexes = np.where(row == v)[0]
            # current label
            __labels = set(y[indexes])

            # Main condition
            if labels is not None and variation > 0:
                # Testing for transition to find the essential cut-points
                if (len(labels) > 1 or len(__labels) > 1) or labels != __labels:
                    # cut-point id
                    cid = len(cutpoints)
                    cutpoints[cid] = (att, u + variation / 2.0)

            labels = __labels
            # previous equals current
            u = v
    new_dict = {}
    # Iterate over the values in the original dictionary
    for key, value in cutpoints.items():
        first_element = value[0]
        second_element = value[1]

        # Check if the first_element is already a key in new_dict
        if first_element in new_dict:
            new_dict[first_element].append(second_element)
        else:
            new_dict[first_element] = [second_element]
    # Generate combinations of the second elements within each group
    for key, value in new_dict.items():
        comb = combinations(value, 2)
        # Append the combinations to the value list
        for c in comb:
            new_dict[key].append(c)
    arrays = []
    for attr, cutpoints in new_dict.items():
        for cutpoint in cutpoints:
            row = X.T[attr]
            if isinstance(cutpoint, tuple):
                lowerbound = cutpoint[0] <= row.reshape(X.shape[0], 1)
                upperbound = row.reshape(X.shape[0], 1) < cutpoint[1]
                row = np.logical_and(lowerbound, upperbound)
                arrays.append(row)
            else:
                row = row.reshape(X.shape[0], 1) >= cutpoint
                arrays.append(row)
    
    Xbin = np.concatenate(arrays, axis=1)

    bin_attr_dim = [i for i in range(1, Xbin.shape[1] + 1)]
    df = pd.DataFrame(Xbin, columns=bin_attr_dim)
    start = 0
    dict_parent_children = {}
    for key, list_value in new_dict.items():
        dict_parent_children[key] = list(df.columns[start: start + len(list_value)])
        start += len(list_value)

    return Xbin, df, dict_parent_children

测试代码示例：

# 小数据集测试
X, y = datasets.load_iris(return_X_y=True)
bin_dataset, data, dict_parent_children = binarize_dataset(X, y)
con_matrix = matrix(bin_dataset, y)

# 较大数据集测试
X, y = datasets.load_breast_cancer(return_X_y=True)
bin_dataset, data, dict_parent_children = binarize_dataset(X, y)
con_matrix = matrix(bin_dataset, y)

1. 用Numpy广播替代多层嵌套循环 #

原函数的四层循环核心是对不同分组的样本两两执行异或操作，利用Numpy的广播机制可以直接批量计算，彻底避免Python循环的性能开销。

优化后的函数：

import numpy as np
import pandas as pd
import time

def matrix_optimized(Xbin, y):
    labels = np.unique(y)
    con_matrix = []
    start = time.time()
    
    for i in range(len(labels)):
        group_i = Xbin[y == labels[i]]  # 形状: (样本数i, 属性维度)
        for j in range(i + 1, len(labels)):
            group_j = Xbin[y == labels[j]]  # 形状: (样本数j, 属性维度)
            # 通过增加维度实现广播：group_i变为(样本数i, 1, 属性维度)，group_j变为(1, 样本数j, 属性维度)
            # 异或后得到(样本数i, 样本数j, 属性维度)，再展平为(样本数i*样本数j, 属性维度)
            xor_result = np.bitwise_xor(group_i[:, np.newaxis], group_j).reshape(-1, Xbin.shape[1])
            con_matrix.append(xor_result)
    
    constraint_matrix = np.vstack(con_matrix)
    end = time.time()
    print("优化版本耗时: ", end - start)
    
    bin_attr_dim = [i for i in range(1, Xbin.shape[1] + 1)]
    df = pd.DataFrame(constraint_matrix, columns=bin_attr_dim)
    return df

效果：将四层循环压缩为两层，计算逻辑交给Numpy的底层C实现，速度通常能提升10~100倍，具体取决于数据集规模。

2. 用Numba JIT编译加速（进阶优化） #

如果广播后的内存占用过高，或者想进一步压榨性能，可以用Numba对循环进行即时编译，让Python循环接近原生C代码的速度。

from numba import jit

@jit(nopython=True)
def compute_xor_batch(group_i, group_j, feat_dim):
    result = np.empty((group_i.shape[0] * group_j.shape[0], feat_dim), dtype=np.bool_)
    idx = 0
    for u in group_i:
        for v in group_j:
            result[idx] = np.bitwise_xor(u, v)
            idx += 1
    return result

def matrix_numba(Xbin, y):
    labels = np.unique(y)
    con_matrix = []
    start = time.time()
    
    for i in range(len(labels)):
        group_i = Xbin[y == labels[i]]
        for j in range(i + 1, len(labels)):
            group_j = Xbin[y == labels[j]]
            xor_result = compute_xor_batch(group_i, group_j, Xbin.shape[1])
            con_matrix.append(xor_result)
    
    constraint_matrix = np.vstack(con_matrix)
    end = time.time()
    print("Numba版本耗时: ", end - start)
    
    bin_attr_dim = [i for i in range(1, Xbin.shape[1] + 1)]
    df = pd.DataFrame(constraint_matrix, columns=bin_attr_dim)
    return df

注意：Numba的nopython模式要求函数内尽量使用NumPy操作，避免Python对象的交互，首次运行会有编译耗时，后续调用会直接复用编译后的代码。

3. 分块计算优化（超大数据集适配） #

如果两组样本数量相乘后结果过大，直接广播会导致内存溢出，可以采用分块计算的方式，每次只处理一部分样本：

def matrix_chunked(Xbin, y, chunk_size=100):
    labels = np.unique(y)
    con_matrix = []
    start = time.time()
    
    for i in range(len(labels)):
        group_i = Xbin[y == labels[i]]
        n_i = group_i.shape[0]
        for j in range(i + 1, len(labels)):
            group_j = Xbin[y == labels[j]]
            # 分块处理group_i，避免一次性占用过多内存
            for chunk_start in range(0, n_i, chunk_size):
                chunk_i = group_i[chunk_start:chunk_start+chunk_size]
                # 对当前块和group_j执行广播异或
                xor_chunk = np.bitwise_xor(chunk_i[:, np.newaxis], group_j).reshape(-1, Xbin.shape[1])
                con_matrix.append(xor_chunk)
    
    constraint_matrix = np.vstack(con_matrix)
    end = time.time()
    print("分块版本耗时: ", end - start)
    
    bin_attr_dim = [i for i in range(1, Xbin.shape[1] + 1)]
    df = pd.DataFrame(constraint_matrix, columns=bin_attr_dim)
    return df

原理：把大的样本组拆成小块，每次只计算一小块与另一组的异或结果，平衡计算速度和内存占用，适合样本量特别大的场景。

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