多维嵌套列表理解及其python,numpy等价物

时间:2014-01-09 15:08:09

标签: python numpy multidimensional-array list-comprehension itertools

我最近发现我可以使用嵌套列表推导来评估参数组合的表达式并将结果保存为数组:

M = np.array(

[ [ [ expression
    for c in C ]
  for b in B ]
for a in A ]

)

其中A,B,C是要迭代的参数,expression是使用a,b,c评估的表达式(例如a+b+c

M将是一个维度为(len(A),len(B),len(C))

的数组

有没有一种简单的方法可以在不使用嵌套列表推导的情况下完成此任务?

我正在寻找一种更加不言自明的解决方案,因为对于在代码中遇到它的人来说,嵌套列表理解可能不是非常易读且容易混淆。再举一个嵌套列表推导的例子

M = np.array(
[
expression
for c in C
    for b in B
        for a in A
]
)

这里,M将是一个平坦的len(A)*len(B)*len(C)维数组。这可以“简化”为:

import itertools as it
M = np.array( [expression for a,b,c in it.product(A,B,C)] )

上述问题是否有类似的解决方案?

感谢。

以下是输入和表达式的具体示例。 P是评估参数a,b或c的函数。在这里,它实际上并没有做任何有用的事情。

A, B, C = np.arange(3), np.arange(4), np.arange(5)
P = lambda *args: sum(args[0][args[1]]*((len(args)>1) or rand()))

M = np.array(

[ [ [ mean( [P(X, a, b) > P(X, c) for X in rand(10,10,10)])
    for c in C ]
  for b in B ]
for a in A ]

)

1 个答案:

答案 0 :(得分:2)

通常,是的,您可以通过将A,B和C转换为numpy数组然后使用广播一步完成来完成此操作。例如

import numpy as np
A, B, C = np.arange(3), np.arange(4), np.arange(5)
result = A[..., None, None] +  B[None, ..., None] + C[None, None, ...]

在这种情况下给我

array([[[0, 1, 2, 3, 4],
        [1, 2, 3, 4, 5],
        [2, 3, 4, 5, 6],
        [3, 4, 5, 6, 7]],

       [[1, 2, 3, 4, 5],
        [2, 3, 4, 5, 6],
        [3, 4, 5, 6, 7],
        [4, 5, 6, 7, 8]],

       [[2, 3, 4, 5, 6],
        [3, 4, 5, 6, 7],
        [4, 5, 6, 7, 8],
        [5, 6, 7, 8, 9]]])

如果你能给出一个具体的例子(输入,预期输出和例子表达),那么我可以使答案更具体。

有关广播的更多内容: