tga.py 文件源码

def _trimmed_mean_1d(arr, k):
    """Calculate trimmed mean on a 1d array.

    Trim values largest than the k'th largest value or smaller than the k'th
    smallest value

    Parameters
    ----------
    arr: ndarray, shape (n,)
        The one-dimensional input array to perform trimmed mean on

    k: int
        The thresholding order for trimmed mean

    Returns
    -------
    trimmed_mean: float
        The trimmed mean calculated
    """
    kth_smallest = np.partition(arr, k)[k-1]
    kth_largest = -np.partition(-arr, k)[k-1]

    cnt = 0
    summation = 0.0
    for elem in arr:
        if elem >= kth_smallest and elem <= kth_largest:
            cnt += 1
            summation += elem
    return summation / cnt