如何有效地创建嵌套循环并附加所有循环的输出（作为多维数组）？

Question

我正在寻找一种在Python 3.x “高效”的方式来创建各种嵌套循环，然后附加每个内部循环（多维数组）的结果。

例如，函数model_A()具有3个参数（ a ， b ， c ），我想列举测试模型的所有可能性。 随意的方式是：

result_a = []
for a_value in a:
    result_a_b = []
    for b_value in b:
        result_a_b_c = []
        for c_value in c:
            result = model_A(a, b, c)
        result_a_b_c.append(result)
    result_a_b.append(result_a_b_c)
result_a.append(result_a_b)

我认为应该有一种“高效”创建嵌套循环并附加结果的方法，而不必在每个循环前创建一个空列表，并将结果附加在每个内部循环的末尾。

Answer 1

大概itertools.product（）具有您所需要的

例如：

a = [1, 2]
b = [3, 4]
c = [5, 6]

for element in itertools.product(a,b,c):
    print(element)

结果是：

(1, 3, 5)
(1, 3, 6)
(1, 4, 5)
(1, 4, 6)
(2, 3, 5)
(2, 3, 6)
(2, 4, 5)
(2, 4, 6)

这可以使您避免深层嵌套的for循环。

Answer 2

使用列表理解可以带来一个小的改进（即使我认为在这种情况下它不那么可读）：

A_LIST = list(range(10))
B_LIST = list(range(10))
C_LIST = list(range(10))

def model_A(a, b, c):
    # mock implementation that does some calculation
    return a + b * c

def f1():
    # your version
    result_a = []
    for a_value in A_LIST:
        result_a_b = []
        for b_value in B_LIST:
            result_a_b_c = []
            for c_value in C_LIST:
                result = model_A(a_value, b_value, c_value)
                result_a_b_c.append(result)
            result_a_b.append(result_a_b_c)
        result_a.append(result_a_b)

    return result_a

def f2():
    # some improvement
    return [
        [
            [model_A(a_value, b_value, c_value)
             for c_value in C_LIST]
            for b_value in B_LIST]
        for a_value in A_LIST]

比较这两个版本，列表理解会稍微快一些：

>>> f1() == f2()
True

>>> import timeit
>>> timeit.timeit('f()', 'from __main__ import f1 as f', number=10000)
1.905
>>> timeit.timeit('f()', 'from __main__ import f2 as f', number=10000)
1.503