python - 根据布尔方阵对 numpy 数组进行排序

Question

假设一个具有 shape 的 numpy 数组A和一个具有 shape的(n,)布尔 numpy 矩阵。B(n,n)

如果B[i][j]是True，那么A[i]应该排到一个位置之前A[j]。

如果B[i][j]是False，那么A[i]应该排序到一个位置之后A[j]。

这些规则仅适用于B[i][j]低于主对角线的情况。主对角线上或主对角线上方的元素应被忽略。

话虽如此，A根据矩阵进行排序的最有效方法是B什么？

我知道有几种简单的方法可以做到这一点，但我必须执行这个操作数千次，所以我正在寻找一种计算效率高的方法来实现它（可读性不是我主要关心的问题）。

标签： pythonalgorithmnumpysortingnumba

Answer 1

我认为这可以工作，我试图使它成为纯 python3。如果你使用 numpy，它可以做得更简单。我使用了 https://stackoverflow.com/a/57003713/3895321。

from functools import cmp_to_key

B = [[True, False, True],
     [True, True, True],
     [False, False, True], ]

A = [2, 1, 3]

tmp = list(range(len(A)))

def compare(i, j):
    if B[i][j]:
        return -1
    else:
        return 1

tmp = sorted(tmp, key=cmp_to_key(compare))

sorted_A = [A[i] for i in tmp]
print(tmp)
print(sorted_A)