numpy - 為其右側第一個滿足特定條件的元素創建索引數組

Question

我想在 numpy 中執行以下操作，而不會產生 for 循環：

假設我有一個數組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

您可以使用以下代碼：

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)為假，則該值將與 idx+1 相同。 如果condition(idx+1)為真，則值為 idx+1
2. 如果idx=len(arr)-1則該值將為len(arr)因為該 idx 的右側沒有任何值

Answer 2

這是我自己的答案：

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