Python 替换嵌套 for 循环

Question

我有以下代码:

import itertools
import numpy as np

t = np.random.rand(3,3,3) 

def foo(T):    

    res = np.zeros((3,3,3))
    
    for a in range(3):
        for b in range(3):
            for c in range(3):
                idx = [a,b,c]
                combinations = list(itertools.permutations(idx, len(idx)))
                idx_arrays = tuple(np.array(idx) for idx in zip(*combinations))     
                res[a, b, c] = (1.0/len(combinations))*np.sum(T[idx_arrays])
        
        
    return res

sol = foo(t)

上面的代码应该相当于:

for a in range(3):
        for b in range(3):
            for c in range(3):
                res2[a,b,c] = (1.0/6)*(
                                        t[a,b,c]
                                        + t[a,c,b]
                                        + t[b,a,c]
                                        + t[b,c,a]
                                        + t[c, a, b]
                                        + t[c,b,a]
                                        )

为了使代码更快，我想替换外部的for循环。这可以做到吗？理想情况下，这应该在不使用 numpy 的 einsum 方法的情况下实现。

Answer 1

您可以通过不进行重复计算来开始改进:

def foo(t):
    res = np.zeros_like(t).astype(t.dtype)
    A, B, C = t.shape
    for a in range(A):
        for b in range(a, B):
            for c in range(b, C):
                r = (1.0 / 6) * (
                    t[a, b, c]
                    + t[a, c, b]
                    + t[b, a, c]
                    + t[b, c, a]
                    + t[c, a, b]
                    + t[c, b, a]
                )
                res[a, b, c] = r
                res[a, c, b] = r
                res[b, a, c] = r
                res[b, c, a] = r
                res[c, a, b] = r
                res[c, b, a] = r
    return res