In essence, I need to do a stable partition: All elements of the 1D np.array a that are present in b should be moved to the front of a but apart from that, the relative order of the elements in a should be preserved.
This code seems to do what I want:
from __future__ import print_function
import numpy as np
if __name__=='__main__':
a = np.array([1, 2, 3, 4, 5], dtype=np.int32)
b = np.array([4, 2, 6], dtype=np.int32)
pos = np.in1d(a, b, assume_unique=True)
cnt = np.count_nonzero(pos)
tmp = np.empty(a.size, dtype=np.int32)
tmp[0:cnt] = a[pos]
tmp[cnt: ] = a[np.invert(pos)]
a = tmp
print('after stable partition:', a)
Output:
after stable partition: [2 4 1 3 5]
Is there a cleaner / more efficient way of doing this in NumPy?
