4
\$\begingroup\$

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?

\$\endgroup\$
2
\$\begingroup\$

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]))
| improve this answer | |
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.