Эффективный способ ограничить числовой логический селектор первыми несколькими истинными значениями - PullRequest
Новая
       24

Эффективный способ ограничить числовой логический селектор первыми несколькими истинными значениями

2 голосов
/ 11 июля 2019

У меня есть массив логических selector, который я могу применить к массиву a. (на самом деле не случайно в проблемной области, это просто удобно для примера) . Но я на самом деле хочу выбрать, используя только первые n истинных записей selector (до n = 3 в примере) . Итак, учитывая selector плюс параметр n, как мне сгенерировать select_first_few, используя пустые операции, избегая, таким образом, итеративного цикла? 

>>> import numpy as np
>>> selector = np.random.random(10) > 0.5
>>> a = np.arange(10)
>>> selector
array([ True, False,  True,  True,  True, False,  True, False,  True,
       False])
>>> chosen, others = a[selector], a[~selector]
>>> chosen
array([0, 2, 3, 4, 6, 8])
>>> others
array([1, 5, 7, 9])
>>> select_first_few = np.array([ True, False,  True,  True,  False, False,  False, False,  False,
...        False])
>>> chosen_few, tough_luck = a[select_first_few], a[~select_first_few]
>>> chosen_few
array([0, 2, 3])
>>> tough_luck
array([1, 4, 5, 6, 7, 8, 9])

Ответы [ 2 ]

1 голос
/ 11 июля 2019

Подход № 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
1 голос
/ 11 июля 2019

Получить все выбранные индексы в списке и нарезать этот список. Затем воспользуйтесь списком, чтобы получить данные по выбранным индексам. 

import numpy as np
selector = np.random.random(10) > 0.5
data = np.arange(10)

choosen_indices = np.where(selector)

#select first 3 choosen
choosen_few_indices = choosen_indices[:3]
choosen_few = [data[i] for i in choosen_few_indices]

# if you are also interested in the not choosen data
not_choosen_indices = list(set(range(len(data))) - set(choosen_indices))
# proceed ...
...