Python实现两个数组对应位置所有可能组合生成矩阵的方法

需求说明 #

给定两个等长数组，需要生成所有对应位置二选一的组合，构成最终矩阵。例如输入array1 = ['A', 'B', 'C']、array2 = ['a', 'b', 'c']时，最终会得到8种组合的2维矩阵。

实现代码 #

方法1：二进制掩码实现（无需依赖第三方库） #

def generate_comb_matrix(arr1, arr2):
    if len(arr1) != len(arr2):
        raise ValueError("两个输入数组长度必须一致")
    n = len(arr1)
    matrix = []
    # 掩码从0到2^n -1，每一位代表对应位置选第一个还是第二个数组的元素
    for mask in range(1 << n):
        row = []
        for idx in range(n):
            row.append(arr2[idx] if mask & (1 << idx) else arr1[idx])
        matrix.append(row)
    return matrix

# 测试调用
array1 = ['A', 'B', 'C']
array2 = ['a', 'b', 'c']
res = generate_comb_matrix(array1, array2)
for row in res:
    print(row)

方法2：使用itertools.product简化实现 #

from itertools import product

array1 = ['A', 'B', 'C']
array2 = ['a', 'b', 'c']
# 先把两个数组按位置打包为可选元素对，再求笛卡尔积
res = [list(comb) for comb in product(*zip(array1, array2))]

输出结果 #

两种方法得到的输出一致，如下：

['A', 'B', 'C']
['a', 'B', 'C']
['A', 'b', 'C']
['a', 'b', 'C']
['A', 'B', 'c']
['a', 'B', 'c']
['A', 'b', 'c']
['a', 'b', 'c']

注：你给出的期望输出最后一行存在笔误，全组合共8种不同的行，最后一行为['a', 'b', 'c']。

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