Stable partition in Numpy

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?

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]))

StackExchange Code Review Q#60805, answer score: 2