python排序返回索引_python – 交叉并排序两个numpy数组的索引

这是一个基于intersect1d实现的选项,这是相当简单的.它需要一次调用argsort.

公认的简单测试通过了.

import numpy as np

def my_intersect(x, y):

"""my_intersect(x, y) -> xm, ym

x, y: 1-d arrays of unique values

xm, ym: indices into x and y giving sorted intersection

"""

# basic idea taken from numpy.lib.arraysetops.intersect1d

aux = np.concatenate((x, y))

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

ym = xym[xym >= len(x)] - len(x)

return xm, ym

def check_my_intersect(x, y):

mx, my = my_intersect(x, y)

assert (x[mx] == np.intersect1d(x, y)).all()

# not really necessary: np.intersect1d returns a sorted list

assert (x[mx] == sorted(x[mx])).all()

assert (x[mx] == y[my]).all()

def random_unique_unsorted(n):

while True:

x = np.unique(np.random.randint(2*n, size=n))

if len(x):

break

np.random.shuffle(x)

return x

x = np.array([4, 1, 10, 5, 8, 13, 11])

y = np.array([20, 5, 4, 9, 11, 7, 25])

check_my_intersect(x, y)

for i in range(20):

x = random_unique_unsorted(100+i)

y = random_unique_unsorted(200+i)

check_my_intersect(x, y)

编辑：“注意”评论令人困惑(使用…作为语音省略号,忘了它也是一个Python运算符).