加權中位數

問題描述為： 一個無序的數列，每個數有其對應的權重，權重為非負整數，代表數列中的數字出現的次數。要求找出這一無序數列中的中位數。

 

1. 直接解法，先對該數列和權重排序。然后找出累計權重為中位數的數字。 時間復雜度為排序的 O（nlog（n）+n）

import numpy as np

def weighted_median(data, weights):
    """
    Args:
      data (list or numpy.array): data
      weights (list or numpy.array): weights
    """
    data, weights = np.array(data).squeeze(), np.array(weights).squeeze()
    s_data, s_weights = map(np.array, zip(*sorted(zip(data, weights))))
    midpoint = 0.5 * sum(s_weights)
    if any(weights > midpoint):
        w_median = (data[weights == np.max(weights)])[0]
    else:
        cs_weights = np.cumsum(s_weights)
        idx = np.where(cs_weights <= midpoint)[0][-1]
        if cs_weights[idx] == midpoint:
            w_median = np.mean(s_data[idx:idx+2])
        else:
            w_median = s_data[idx+1]
    return w_median

def test_weighted_median():
    print("hello, world")
    data = [
        [7, 1, 2, 4, 10],
        [7, 1, 2, 4, 10],
        [7, 1, 2, 4, 10, 15],
        [1, 2, 4, 7, 10, 15],
        [0, 10, 20, 30],
        [1, 2, 3, 4, 5],
        [30, 40, 50, 60, 35],
        [2, 0.6, 1.3, 0.3, 0.3, 1.7, 0.7, 1.7, 0.4],
    ]
    weights = [
        [1, 1/3, 1/3, 1/3, 1],
        [1, 1, 1, 1, 1],
        [1, 1/3, 1/3, 1/3, 1, 1],
        [1/3, 1/3, 1/3, 1, 1, 1],
        [30, 191, 9, 0],
        [10, 1, 1, 1, 9],
        [1, 3, 5, 4, 2],
        [2, 2, 0, 1, 2, 2, 1, 6, 0],
    ]
    answers = [7, 4, 8.5, 8.5, 10, 2.5, 50, 1.7]
    for datum, weight, answer in zip(data, weights, answers):
        assert(weighted_median(datum, weight) == answer)

if __name__ == "__main__":
    test_weighted_median()

 

 2. 按照快速排序的思路，先找到一個數字，然后 按照該數字將數列划分成左右兩段，根據左右兩段的權重之和，遞歸調用左半側或者右半側數列。