Jak wymienić tylko pierwsze n elementów w tablicy numpy, które są większe niż określona wartość?

Question

Mam tablicę myA takiego:

array([ 7, 4, 5, 8, 3, 10]) 

Jeśli chcę zastąpić wszystkie wartości, które są większe niż wartość val przez 0, mogę po prostu zrobić:

myA[myA > val] = 0 

który daje mi sygnał wyjściowy (dla val = 5):

array([0, 4, 5, 0, 3, 0]) 

jednak moim celem jest zastąpienie nie wszystkich, ale tylko pierwszy n elementy tej tablicy, które są większe niż wartość val.

Tak więc, jeśli n = 2 mój pożądany rezultat będzie wyglądać następująco (10 jest trzeci element, a zatem nie powinny zostały zastąpione):

array([ 0, 4, 5, 0, 3, 10]) 

Prosta implementacja będzie:

import numpy as np 

myA = np.array([7, 4, 5, 8, 3, 10]) 
n = 2 
val = 5 

# track the number of replacements 
repl = 0 

for ind, vali in enumerate(myA): 

    if vali > val: 

     myA[ind] = 0 
     repl += 1 

     if repl == n: 
      break 

Że działa, ale może ktoś może wymyślić mądry sposób maskowania !?

Answer 1

Poniższy powinno działać:

myA[(myA > val).nonzero()[0][:2]] = 0 

od nonzero zwróci indeksy, w których tablica boolowska myA > val jest niezerowa, np. True.

Na przykład:

In [1]: myA = array([ 7, 4, 5, 8, 3, 10]) 

In [2]: myA[(myA > 5).nonzero()[0][:2]] = 0 

In [3]: myA 
Out[3]: array([ 0, 4, 5, 0, 3, 10]) 

Answer 2

Można użyć tego samego rozwiązania jak here z konwersji Ci np.array do pd.Series:

s = pd.Series([ 7, 4, 5, 8, 3, 10]) 
n = 2 
m = 5 
s[s[s>m].iloc[:n].index] = 0 

In [416]: s 
Out[416]: 
0  0 
1  4 
2  5 
3  0 
4  3 
5 10 
dtype: int64 

Krok po kroku wyjaśnienia:

In [426]: s > m 
Out[426]: 
0  True 
1 False 
2 False 
3  True 
4 False 
5  True 
dtype: bool 

In [428]: s[s>m].iloc[:n] 
Out[428]: 
0 7 
3 8 
dtype: int64 

In [429]: s[s>m].iloc[:n].index 
Out[429]: Int64Index([0, 3], dtype='int64') 

In [430]: s[s[s>m].iloc[:n].index] 
Out[430]: 
0 7 
3 8 
dtype: int64 

Wyjście w In[430] wygląda tak samo jak In[428] ale w 428 jest to kopię i 430 oryginalnych serii.

Jeśli musisz np.array można użyć values metody:

In [418]: s.values 
Out[418]: array([ 0, 4, 5, 0, 3, 10], dtype=int64) 

Answer 3

Ostateczne rozwiązanie jest bardzo proste:

import numpy as np 
myA = np.array([7, 4, 5, 8, 3, 10]) 
n = 2 
val = 5 

myA[np.where(myA > val)[0][:n]] = 0 

print(myA) 

wyjściowa:

[ 0 4 5 0 3 10] 

Answer 4

Oto inna możliwość (testowane), prawdopodobnie nie lepiej niż nonzero:

def truncate_mask(m, stop): 
    m = m.astype(bool, copy=False) # if we allow non-bool m, the next line becomes nonsense 
    return m & (np.cumsum(m) <= stop) 

myA[truncate_mask(myA > val, n)] = 0 

Unikając budynku i korzystania wyraźny wskaźnik może skończyć się z nieco lepszej wydajności ... ale musisz to przetestować, żeby się przekonać.

Edit 1: gdy jesteśmy na temat możliwości, można też spróbować:

def truncate_mask(m, stop): 
    m = m.astype(bool, copy=True) # note we need to copy m here to safely modify it 
    m[np.searchsorted(np.cumsum(m), stop):] = 0 
    return m 

Edycja 2 (następnego dnia): Właśnie przetestowane i wydaje że cumsum jest faktycznie gorszy niż nonzero, przynajmniej z kinds of values, którego używałem (więc żadne z powyższych podejść nie jest warte użycia). Z ciekawości, ja też próbowałem z Numba:

import numba 

@numba.jit 
def set_first_n_gt_thresh(a, val, thresh, n): 
    ii = 0 
    while n>0 and ii < len(a): 
     if a[ii] > thresh: 
      a[ii] = val 
      n -= 1 
     ii += 1 

This tylko iteracje nad tablicy raz, czy raczej tylko iteracje nad niezbędną część tablicy raz, nawet nie dotykając ostatnią część. Daje to znacznie lepszą wydajność dla małych n, ale nawet w najgorszym przypadku n>=len(a) podejście to jest szybsze.