NumPy: function for simultaneous max() and min()

后端 未结 12 923
天涯浪人
天涯浪人 2020-11-27 04:54

numpy.amax() will find the max value in an array, and numpy.amin() does the same for the min value. If I want to find both max and min, I have to call both functions, which

12条回答
  •  北荒
    北荒 (楼主)
    2020-11-27 05:15

    Just to get some ideas on the numbers one could expect, given the following approaches:

    import numpy as np
    
    
    def extrema_np(arr):
        return np.max(arr), np.min(arr)
    
    import numba as nb
    
    
    @nb.jit(nopython=True)
    def extrema_loop_nb(arr):
        n = arr.size
        max_val = min_val = arr[0]
        for i in range(1, n):
            item = arr[i]
            if item > max_val:
                max_val = item
            elif item < min_val:
                min_val = item
        return max_val, min_val
    
    import numba as nb
    
    
    @nb.jit(nopython=True)
    def extrema_while_nb(arr):
        n = arr.size
        odd = n % 2
        if not odd:
            n -= 1
        max_val = min_val = arr[0]
        i = 1
        while i < n:
            x = arr[i]
            y = arr[i + 1]
            if x > y:
                x, y = y, x
            min_val = min(x, min_val)
            max_val = max(y, max_val)
            i += 2
        if not odd:
            x = arr[n]
            min_val = min(x, min_val)
            max_val = max(x, max_val)
        return max_val, min_val
    
    %%cython -c-O3 -c-march=native -a
    #cython: language_level=3, boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True, infer_types=True
    
    
    import numpy as np
    
    
    cdef void _extrema_loop_cy(
            long[:] arr,
            size_t n,
            long[:] result):
        cdef size_t i
        cdef long item, max_val, min_val
        max_val = arr[0]
        min_val = arr[0]
        for i in range(1, n):
            item = arr[i]
            if item > max_val:
                max_val = item
            elif item < min_val:
                min_val = item
        result[0] = max_val
        result[1] = min_val
    
    
    def extrema_loop_cy(arr):
        result = np.zeros(2, dtype=arr.dtype)
        _extrema_loop_cy(arr, arr.size, result)
        return result[0], result[1]
    
    %%cython -c-O3 -c-march=native -a
    #cython: language_level=3, boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True, infer_types=True
    
    
    import numpy as np
    
    
    cdef void _extrema_while_cy(
            long[:] arr,
            size_t n,
            long[:] result):
        cdef size_t i, odd
        cdef long x, y, max_val, min_val
        max_val = arr[0]
        min_val = arr[0]
        odd = n % 2
        if not odd:
            n -= 1
        max_val = min_val = arr[0]
        i = 1
        while i < n:
            x = arr[i]
            y = arr[i + 1]
            if x > y:
                x, y = y, x
            min_val = min(x, min_val)
            max_val = max(y, max_val)
            i += 2
        if not odd:
            x = arr[n]
            min_val = min(x, min_val)
            max_val = max(x, max_val)
        result[0] = max_val
        result[1] = min_val
    
    
    def extrema_while_cy(arr):
        result = np.zeros(2, dtype=arr.dtype)
        _extrema_while_cy(arr, arr.size, result)
        return result[0], result[1]
    

    (the extrema_loop_*() approaches are similar to what is proposed here, while extrema_while_*() approaches are based on the code from here)

    The following timings:

    indicate that the extrema_while_*() are the fastest, with extrema_while_nb() being fastest. In any case, also the extrema_loop_nb() and extrema_loop_cy() solutions do outperform the NumPy-only approach (using np.max() and np.min() separately).

    Finally, note that none of these is as flexible as np.min()/np.max() (in terms of n-dim support, axis parameter, etc.).

    (full code is available here)

提交回复
热议问题