In essence, I need to do a stable partition: All elements of the 1D
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)
after stable partition: [2 4 1 3 5]
Is there a cleaner / more efficient way of doing this in NumPy?