根據布爾python獲得數組部分的中位數

Question

我有兩個相同長度的數組，第一個是布爾數組，第二個包含相應的值。

flag = [0,0,0,1,1,0,0,0,1,1,1,1,0,1,1] 
values = [1,5,6,8,5,6,2,0,1,9,3,8,3,6,2] 

欲返回包含對應於布爾矩陣的1各部分的中值的中位數的陣列。

例如

flag = [0,0,0,1, 1, 0,0,0, 1, 1, 1, 1, 0,1,1] 
result = [0,0,0,6.5,6.5,0,0,0,5.5,5.5,5.5,5.5,0,4,4] 

我unesthetic的做法是要做到：

result = np.zeros(values.shape[0]) 
vect = [] 
idx = [] 
for n in np.arange(result.size): 
    if flag[n] > 0: 
     vect.append(values[n]) 
     idx.append(n) 
    elif flag[n] == 0: 
     result[idx] = np.median(vect) 
     vect = [] 
     idx = [] 
    result[idx] = np.median(vect) 

它運作良好，但它不是很Python的很慢，因爲我有非常大的數組。

Answer 1

我們可以使用np.diff找到0和1之間的轉換。然後循環0/1和1/0轉換對，並從中間值中取中間值。

結果循環遍歷每一組。

flag = [0,0,0,1,1,0,0,0,1,1,1,1,0,1,1] 
values = [1,5,6,8,5,6,2,0,1,9,3,8,3,6,2] 

d = np.diff(np.concatenate([[0], flag, [0]])) # Add and append a 0 so the procedure also works if flags start or end with 1. 

begin = np.flatnonzero(d==1) 
end = np.flatnonzero(d==-1) 

result = np.zeros_like(values, dtype=float) 

for a, b in zip(begin, end): 
    result[a:b] = np.median(values[a:b]) 

print(result) 
# [ 0. 0. 0. 6.5 6.5 0. 0. 0. 5.5 5.5 5.5 5.5 0. 4. 4. ]