如何從 numpy 數組中提取居中的 window？

Question

假設我有以下 numpy 數組：

>>> a = np.arange(0,21,1)
>>> a
array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20])

現在假設我想選擇一個長度為N的 window ，其中2 < N <= a.shape[0] ，這樣 window 圍繞數組a元素“居中”。 例如，如果我想將長度為N = 5的 window 圍繞數組a中的元素10居中，那么這個 window 將是：

>>> idx = 10 # index of the array element 10
>>> N = 5 # window size
>>> a[idx - N//2:idx + N//2 + 1]
array([ 8,  9, 10, 11, 12])

此方法適用於不在陣列邊緣附近的 windows，但我無法使其正常工作。 例如，如果我想在a中的元素2周圍提取長度為N = 7的 window ，那么我得到的是：

>>> idx = 2
>>> N = 7
>>> a[idx - N//2:idx + N//2 + 1]
array([], dtype=int32)

但是我想要的是：

>>> a[0:7]
array([0, 1, 2, 3, 4, 5, 6])

我如何將這種方法推廣到 windows a邊緣附近？

Answer 1

嘗試：

idx = 2
start = min(idx - N//2, 0)
a[start:start + N]

請注意，這不是以idx=2為中心。

Answer 2

根據Quang Hoang 的回答，以下是有效的：

import numpy as np

a = np.arange(0,21,1)
idx = 5 # desired element index
N = 7 # window length

if N % 2: # if window length is odd
    step = N // 2
else: # if window length is even
    step = int(N/2 - 1)

# make sure starting index is between 0 and a.shape[0] - N
start = min(max(idx-step,0),a.shape[0] - N)
window = a[start:start + N]