根据bool 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]

我不美观的做法是：

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)

它工作得很好，但它不是很pythonic而且很慢，因为我使用的是非常大的数组。

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. ]