Question

如何进行阵列切片？给定2,3; 7,8和a，我想获得数组start_idx中的项目end_idx。我尝试了a[[0,1], start_idx: end_idx]，但失败了。

import numpy as np
a = np.array([[1,2,3,4], [5,6,7,8]])
start_idx=[1,2]
end_idx=[3,4]
a[[0,1], start_idx: end_idx]] <- failed
[a[0,1:3], a[1,2:4]] #[array([2, 3]), array([7, 8])] <- wanted result

Answer 1

您可以对enumerate和zip使用列表推导：

res = [a[idx, start: end] for idx, (start, end) in enumerate(zip(start_idx, end_idx))]

# [array([2, 3]), array([7, 8])]

这类似于@Bazingaa的解决方案，但是直接在start_idx / end_idx上进行迭代，而不是在位置列表索引上进行迭代。

性能优势是微不足道的，并且只有在大型阵列中才会变得明显。但是，您可能会发现枚举逻辑更具可读性。

n = 100000
a = np.array([[1,2,3,4], [5,6,7,8]]*n)
start_idx = [1,2] * n
end_idx = [3,4] * n

def indexing(a, start_idx, end_idx):
    return [a[i, start_idx[i]:end_idx[i]] for i in range(len(start_idx))]

def enumerating(a, start_idx, end_idx):
    return [a[idx, start: end] for idx, (start, end) in enumerate(zip(start_idx, end_idx))]

%timeit indexing(a, start_idx, end_idx)     # 140 ms
%timeit enumerating(a, start_idx, end_idx)  # 124 ms

Answer 2

您可以使用列表理解：

result = [a[i, start_idx[i]:end_idx[i]] for i in range(len(start_idx))]
print (result)

输出

[array([2, 3]), array([7, 8])]

我可以像在python中进行数组索引一样进行数组切片吗？

2 个答案: