intersect_two_numpy_arrays

numpy_utils_intersect.py

import numpy as np

def intersect_indcs_two_arrays(x_array, y_array):
    """intersect_indcs_two_arrays(x, y) -> xm, ym
    x, y: 1-d arrays of unique values
    xm, ym: indices into x and y giving sorted intersection
    code taken from http://stackoverflow.com/questions/33529057/indices-that-intersect-and-sort-two-numpy-arrays?lq=1
    """
    assert np.all(np.logical_not(map(are_there_duplicates, [x_array, y_array]))), "Intersect assumes arrays are unique."
    # basic idea taken from numpy.lib.arraysetops.intersect1d
    aux = np.concatenate((x_array, y_array))
    sidx = aux.argsort()
    # Note: intersect1d uses aux[:-1][aux[1:]==aux[:-1]] here - I don't know why the first [:-1] is necessary
    inidx = aux[sidx[1:]] == aux[sidx[:-1]]

    # quicksort is not stable, so must do some work to extract indices
    # (if stable, sidx[inidx.nonzero()]  would be for x)
    # interlace the two sets of indices, and check against lengths
    xym = np.vstack((sidx[inidx.nonzero()],
                     sidx[1:][inidx.nonzero()])).T.flatten()

    xm = xym[xym < len(x_array)]
    ym = xym[xym >= len(x_array)] - len(x_array)
    return xm, ym