首页 > 解决方案 > 将 numpys 花式索引与切片相结合

问题描述

我目前有一个二维 numpy 数组 shape (m, n)。此外,我有两个(m, p)索引数组i1i2. 索引总是连续的!

import numpy as np
t = np.array([[-1, -1,  0,  0,  1,  2,  2],
              [-1, -1,  0,  1,  2,  3,  3],
              [0,  0,  1,  2,  2,  3,  3]])

i1 = np.array([3, 2, 2])
i2 = np.array([4, 3, 3])

如何使用数组i1i2切片t以获得以下子矩阵?

expected_t = np.array([
                      [0, 1],
                      [0, 1],
                      [1, 2]
])

那是

expected_t[0, :] = t[0, i1[0]:i2[0]]
expected_t[1, :] = t[1, i1[1]:i2[1]]
expected_t[2, :] = t[2, i1[2]:i2[2]]

此外,这是否可以在不通过创建视图复制数据的情况下完成?

提前感谢所有帮助!

标签: pythonnumpymatrix-indexing

解决方案


在 numpy 中使用精美的索引:

t[np.arange(3).reshape(3,1), np.vstack((i1,i2)).T]

或者

t[np.arange(3), np.vstack((i1,i2))].T

两者都会有结果:

array([[0, 1],
   [0, 1],
   [1, 2]])

推荐阅读