Эффективный поиск следующего большего в другом массиве

Question

Можно ли удалить циклы for в этой функции и ускорить процесс? Я не смог получить те же результаты с векторными методами для этой функции. Или есть другой вариант?

import numpy as np

indices = np.array(
    [814, 935, 1057, 3069, 3305, 3800, 4093, 4162, 4449])

within = np.array(
    [193, 207, 243, 251, 273, 286, 405, 427, 696,
     770, 883, 896, 1004, 2014, 2032, 2033, 2046, 2066,
     2079, 2154, 2155, 2156, 2157, 2158, 2159, 2163, 2165,
     2166, 2167, 2183, 2184, 2208, 2210, 2212, 2213, 2221,
     2222, 2223, 2225, 2226, 2227, 2281, 2282, 2338, 2401,
     2611, 2612, 2639, 2640, 2649, 2700, 2775, 2776, 2785,
     3030, 3171, 3191, 3406, 3427, 3527, 3984, 3996, 3997,
     4024, 4323, 4331, 4332])


def get_first_ind_after(indices, within):
    """returns array of the first index after each listed in indices

    indices and within must be sorted ascending
    """
    first_after_leading = []
    for index in indices:

        for w_ind in within:

            if w_ind > index:
                first_after_leading.append(w_ind)

                break

    # convert to np array
    first_after_leading = np.array(first_after_leading).flatten()

    return np.unique(first_after_leading)

Он должен возвращать следующее наибольшее число для каждого в массиве индексов, если оно есть.

# Output:
[ 883 1004 2014 3171 3406 3984 4323]

Answer 1

Вот один на основе np.searchsorted -

def next_greater(indices, within):
    idx = np.searchsorted(within, indices)
    idxv = idx[idx<len(within)]
    idxv_unq = np.unique(idxv)
    return within[idxv_unq]

В качестве альтернативы, idxv_unq может быть вычислено примерно так и должно быть более эффективным -

idxv_unq = idxv[np.r_[True,idxv[:-1] != idxv[1:]]]

Answer 2

Попробуйте это:

[within[within>x][0] if len(within[within>x])>0 else 0 for x in indices]

Как в,

In [35]: import numpy as np
    ...: indices = np.array([814, 935, 1057, 3069, 3305, 3800, 4093, 4162, 4449])
    ...:
    ...: within = np.array(
    ...:     [193, 207, 243, 251, 273, 286, 405, 427, 696,
    ...:      770, 883, 896, 1004, 2014, 2032, 2033, 2046, 2066,
    ...:      2079, 2154, 2155, 2156, 2157, 2158, 2159, 2163, 2165,
    ...:      2166, 2167, 2183, 2184, 2208, 2210, 2212, 2213, 2221,
    ...:      2222, 2223, 2225, 2226, 2227, 2281, 2282, 2338, 2401,
    ...:      2611, 2612, 2639, 2640, 2649, 2700, 2775, 2776, 2785,
    ...:      3030, 3171, 3191, 3406, 3427, 3527, 3984, 3996, 3997,
    ...:      4024, 4323, 4331, 4332])

In [36]: [within[within>x][0] if len(within[within>x])>0 else 0 for x in indices]
Out[36]: [883, 1004, 2014, 3171, 3406, 3984, 4323, 4323, 0]

Это питонический подход, называемый понимание списка это сокращенная версия цикла foreach. Так что, если бы я расширил это:

result = []
for x in indices:
    # This next line is a boolean index into the array, if returns all of the items in the array that have a value greater than x
    y = within[within>x]
    # At this point, y is an array of all the items which are larger than x.  Since you wanted the first of these items, we'll just take the first item off of this new array, but it is possible that y is None (there are no values that match the condition), so there is a check for that
    if len(y) > 0:
         z = y[0]
    else:
         z = 0 # or None or whatever you like
    # Now add this value to the array that we are building
    result.append(z)
# Now result has the array

Я написал это так, потому что он использует векторные операции (то есть логическую маску), а также использует понимание списка, что является гораздо более простым способом написать foreach, который возвращает массив.