Подход № 1
Одним из подходов будет использование cumsum и argmax для получения экстента, а затем среза для установки False
-
In [40]: n = 3
In [41]: selector
Out[41]:
array([ True, False, True, True, True, False, True, False, True,
False])
In [42]: selector[(selector.cumsum()>n).argmax():] = 0
In [43]: selector # your select_first_few mask
Out[43]:
array([ True, False, True, True, False, False, False, False, False,
False])
Затем используйте этот новый selector
для выбора и отмены выбора элементов из входного массива.
Подход № 2
Другим подходом будет маска-маска -
n = 3
C = np.count_nonzero(selector)
newmask = np.zeros(C, dtype=bool)
newmask[:n] = 1
selector[selector] = newmask
Пробный прогон -
In [62]: selector
Out[62]:
array([ True, False, True, True, True, False, True, False, True,
False])
In [63]: n = 3
...: C = np.count_nonzero(selector)
...: newmask = np.zeros(C, dtype=bool)
...: newmask[:n] = 1
...: selector[selector] = newmask
In [64]: selector
Out[64]:
array([ True, False, True, True, False, False, False, False, False,
False])
Или сделать его короче с конкатенацией логических значений на лету -
n = 3
C = np.count_nonzero(selector)
selector[selector] = np.r_[np.ones(n,dtype=bool),np.zeros(C-n,dtype=bool)]
Подход № 3
Самый простой -
selector &= selector.cumsum()<=n