在2D阵列上矢量化NumPy三重产品

Question

我正在尝试在以下名为N x N的{​​{1}}数组上对以下三重产品操作进行矢量化：

p

我认为for j in range(len(p)): for k in range(len(p)): for l in range(len(p)): h[j, k, l] = p[j, k] * p[k, l] * p[l, j] - p[j, l] * p[l, k] * p[k, j] 应该在这里使用，尽管我实际上并没有对重复的指数进行总结，但是我还没能把它固定下来。想法？

Answer 1

只需将这些循环迭代器移植为字符串表示法，我们就会有一个基于einsum的解决方案 -

h = np.einsum('jk,kl,lj->jkl',p,p,p) - np.einsum('jl,lk,kj->jkl',p,p,p)

基本上是一个与扩展相关的问题（因为我们没有减少任何轴），我们也可以简单地使用NumPy broadcasting，在不同的地方引入带有None/np.newaxis的新轴以允许扩展，就像这样 -

h = p[...,None]*p*p[:,None,:].T - p[:,None,:]*p.T*p.T[...,None]