For each label in one array set the first k occurrences to False in another array

user_10582903 0
Viewed 73

I have two (sorted) arrays, A and B, of different lengths each containing unique labels that are repeated a number of times. The count for each label in A is less than or equal to that in B. All labels in A will be in B, but some labels in B do not appear in A.

I need an object the same length as B where, for each label i in A (which occurs k_i times), the first k_i occurrences of label i in B need to be set to False. The remaining elements should be True.

The following code gives me what I need, but if A and B are large, this can take a long time:

import numpy as np

# The labels and their frequency
A = np.array((1,1,2,2,3,4,4,4))
B = np.array((1,1,1,1,1,2,2,3,3,4,4,4,4,4,5,5))

A_uniq, A_count = np.unique(A, return_counts = True)
new_ind = np.ones(B.shape, dtype = bool)
for i in range(len(A_uniq)):
    new_ind[np.where(B == A_uniq[i])[0][:A_count[i]]] = False

print(new_ind)
#[False False  True  True  True False False False  True False False False
#  True  True  True  True]

Is there a faster or more efficient way to do this? I feel like I may be missing some obvious broadcasting or vectorized solution.

numpy python
3 Answers

Here's one with np.searchsorted -

idx = np.searchsorted(B, A_uniq)
id_ar = np.zeros(len(B),dtype=int)
id_ar[idx] = 1
id_ar[A_count+idx] -= 1
out = id_ar.cumsum()==0

We can optimize further to compute A_uniq,A_count using its sorted nature instead of using np.unique, like so -

mask_A = np.r_[True,A[:-1]!=A[1:],True]
A_uniq, A_count = A[mask_A[:-1]], np.diff(np.flatnonzero(mask_A))
user_32938810

Example without numpy

A = [1,1,2,2,3,4,4,4]
B = [1,1,1,1,1,2,2,3,3,4,4,4,4,4,5,5]

a_i = b_i = 0
while a_i < len(A):
  if A[a_i] == B[b_i]:
    a_i += 1
    B[b_i] = False
  else:
    B[b_i] = True
  b_i += 1
# fill the rest of B with True
B[b_i:] = [True] * (len(B) - b_i)
# [False, False, True, True, True, False, False, False, True, False, False, False, True, True, True, True]
user_57500020

This solution is inspired by the one by @Divakar, using itertools.groupby:

import numpy as np
from itertools import groupby
A = np.array((1, 1, 2, 2, 3, 4, 4, 4))
B = np.array((1, 1, 1, 1, 1, 2, 2, 3, 3, 4, 4, 4, 4, 4, 5, 5))

indices = [key + i for key, group in groupby(np.searchsorted(B, A)) for i, _ in enumerate(group)]
result = np.ones_like(B, dtype=np.bool)
result[indices] = False

print(result)

Output

[False False  True  True  True False False False  True False False False
  True  True  True  True]

The idea is to use np.searchsorted to find the insertion position of each element of A, as equal elements will have the same insertion position you have to shift by one each of them, hence the groupby. Then create an array of True and set the values of the indices to False.

If you can use pandas, compute the indices like this:

values = np.searchsorted(B, A)
indices = pd.Series(values).groupby(values).cumcount() + values
user_40015920
Related