Question

一个相当难的短语问题。我需要索引一个数组，以便满足所有a和b的以下条件：

arr[a, b, a, b] = data[b]

这可以通过以下方式实现：

import numpy as np

a = 3
b = 2

arr = np.zeros((a, b, a, b))

data = np.array([10, 20])
rngb = np.arange(b)

for i in range(a):
    arr[i, rngb, i, rngb] += data

print arr

这里的目标是删除python for循环。一个简单的答案是在索引和数据数组上使用tile和repeat，但我相信我有一个更容易解决的问题。任何建议都将非常受欢迎。

编辑：另一种但可能更昂贵的计算方式arr：

diag = np.diag(np.ones(a))
arr = np.einsum('ik,jl->ijkl', diag, np.diag(data))

Answer 1

def meth1(a,b):  # old method
    arr = np.zeros((a,b,a,b))
    data = np.arange(b)+1
    rngb = np.arange(b)
    for i in range(a):
        arr[i, rngb, i, rngb] += data
    return arr

def meth2(a,b):  # revised method
    arr = np.zeros((a, b, a, b))
    data = np.arange(b)+1
    arr.ravel()[::a*b+1] = np.tile(data, a)
    return arr

您可以使用简单的lambda测试两者的相等性：

m = lambda a,b: np.any((meth2(a,b)-meth1(a,b)))

索引4维阵列的偏移对角线

1 个答案: