numpy - создать массив индекса для первого элемента справа от него, который удовлетворяет определенному условию

Question

Я хотел бы сделать следующее в numpy без использования a для l oop:

Предположим, у меня есть массив x, и для каждого элемента в x я бы хотел найти индекс первого элемента справа от него, который удовлетворяет определенному условию. Если справа от него нет такого элемента, верните длину массива.

Например, массив равен x = [1, 3, 5, 2, 4, 4, 3, 1, 5], и условие состоит в том, что элемент равен или больше 4. В этом случае результат должен быть y = [2, 2, 4, 4, 5, 8, 8, 8, 9], поскольку индекс первого элемента равен или больше 4 имеет значение 5 и индекс 2 и др.

Answer 1

Вот мой собственный ответ:

x = [1, 3, 5, 2, 4, 4, 3, 1, 5]
index_condition = np.argwhere(x >= 4) # the selection criterion is elements >= 4
index_count = np.ediff1d(index_condition, to_begin=index_condition[0]) # get the count of each index in the final result as consecutive diff
result = np.repeat(index_condition, index_count) # repeat the index
# also need to pad the result to desired length but is easy

Answer 2

Вы можете использовать следующий код:

def condition(a):
    return a>=4

def generate_array(arr, idx):
    if len(arr)==idx+1:
        return [len(arr)]

    n_arr = generate_array(arr, idx+1)
    if condition(arr[idx+1]):
        return [idx+1] + n_arr

    return [n_arr[0]] + n_arr

Тогда:

x = [1, 3, 5, 2, 4, 4, 3, 1, 5]
print(generate_array(x, 0))
[2, 2, 4, 4, 5, 8, 8, 8, 9]

В основном вы можете использовать следующие факты

1. Для любого индекса idx<len(arr)-1, если condition(idx+1) равно false, то значение будет таким же, как для idx + 1. Если condition(idx+1) равно true, тогда значение будет idx + 1
2. Если idx=len(arr)-1, то значение будет len(arr), поскольку справа от этого idx нет значения