Stable partition in Numpy

Question

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?

Answer 1

The actual creation of the pos mask looks perfect to me. You can streamline the rest of the code into a single-liner as:

a = np.concatenate((a[pos], a[~pos]))