Python：从布尔数组获取最后n个True

Question

I have a boolean array and a want to convert that to an array, where only the last_n_trues True values are still True . A simple example:

>>> boolean_array = [False, False, True, True, True, False, False]
>>> last_n_trues = 2
>>> desired_output = [False, False, False, True, True, False, False]

My approach:

>>> import numpy as np
>>> idxs_of_trues = np.where(boolean_array)[0] 
array([2, 3, 4], dtype=int64)

>>> idxs_of_trues_last_n = idxs_of_trues[-last_n_trues:] 
array([3, 4], dtype=int64)

>>> [x in idxs_of_trues_last_n for x in range(0, len(boolean_array))]
[False, False, False, True, True, False, False]

Is there a faster way to do so? Especially the list comprehension seems pretty complicated to me...

Answer 1

You should just be able to simply use np.where

In [116]: x
Out[116]: array([False, False,  True,  True,  True, False, False], dtype=bool)

In [117]: x[np.where(x)[0][:-2]] = False

In [118]: x
Out[118]: array([False, False, False,  True,  True, False, False], dtype=bool)

This just replaces all True that aren't the last 2 with False

This will only work if x is a np.array , so verify that before you try this.

Answer 2

Approach #1 : Here's one with cumsum -

def keep_lastNTrue_cumsum(a, n):
    c = np.count_nonzero(a) # or a.sum()
    a[c - a.cumsum() >= n] = 0
    return a

Approach #2 : Two more with argpartition -

def keep_lastNTrue_argpartition1(a, n):
    c = np.count_nonzero(a) # or a.sum()
    a[np.sort(np.argpartition(a,-n)[-c:])[:-n]] = 0
    return a

def keep_lastNTrue_argpartition2(a, n):
    c = np.count_nonzero(a) # or a.sum()
    p = np.argpartition(a,-n)[-a.sum():]
    cn = c-n
    idx = np.argpartition(p,cn)
    a[p[idx[:cn]]] = 0
    return a

Approach #3 : Another with a bit more of mask usage -

def keep_lastNTrue_allmask(a, n):
    c = a.sum()
    set_mask = np.ones(c, dtype=bool)
    set_mask[:-n] = False
    a[a] = set_mask
    return a

Sample runs -

In [141]: boolean_array = np.array([False, False, True, True, True, False, False])

In [142]: keep_lastNTrue_cumsum(boolean_array, n=2)
Out[142]: array([False, False, False,  True,  True, False, False])

In [143]: boolean_array = np.array([False, False, True, True, True, False, False])

In [144]: keep_lastNTrue_argpartition1(boolean_array, n=2)
Out[144]: array([False, False, False,  True,  True, False, False])

In [145]: boolean_array = np.array([False, False, True, True, True, False, False])

In [146]: keep_lastNTrue_argpartition2(boolean_array, n=2)
Out[146]: array([False, False, False,  True,  True, False, False])

Answer 3

The fastest way without libraries is going to be to clone the list and iterate through it in reverse:

def foo(bools, last_n_trues):
    result = bools[:]
    count = 0
    for i in range(len(bools) - 1, -1, -1):
        if count < last_n_trues:
            if result[i]:
                count += 1
        else:       
            result[i] = False
    return result