在分配期间修改numpy索引

时间:2016-06-13 03:48:38

标签: arrays numpy indexing

我不确定搜索找到正确优化的正确术语。我想将下面的最后四行代码简化为两行,其中+/- 1的添加分别在分配给plusminus变量时完成。

    # generic params to simulate loop conditions
    position = np.arange(10)
    axis = 2

    # actual code to optimise
    plus = np.asarray(position)
    plus[axis] += 1
    minus = np.asarray(position)
    minus[axis] -= 1

澄清这是一个迭代问题:任何不采用通用positionaxis变量的解决方案都是错误的,即明确以下不是解决方案:

plus = np.asarray([0,1,3,3,4,5,6,7,8,9])
plus = np.asarray(range(axis)+[position[axis]+1]+range(axis+1,len(position)))

1 个答案:

答案 0 :(得分:1)

这是一种使用np.in1d 这四行压缩为的方法 -

mask = np.in1d(np.arange(position.size),axis)
plus, minus = position + mask, position - mask

示例运行

让我们测试一下具有另一个索引position的通用6数组 -

In [60]: position
Out[60]: array([1, 0, 6, 8, 1, 7, 1, 3, 1, 6])

In [61]: axis = 6

In [62]: mask = np.in1d(np.arange(position.size),axis)

In [63]: plus, minus = position + mask, position - mask

In [64]: plus
Out[64]: array([1, 0, 6, 8, 1, 7, 2, 3, 1, 6]) # Change at 6th index

In [65]: minus
Out[65]: array([1, 0, 6, 8, 1, 7, 0, 3, 1, 6]) # Change at 6th index