python numpy 加速 2d 重复搜索

Question

我需要在二维 numpy 数组中查找重复项。结果，我想要一个与输入相同长度的列表，该列表指向相应值的第一次出现。例如数组 [[1, 0, 0], [1, 0, 0], [2, 3, 4]] 有两个相等的元素 0 和 1。该方法应该返回 [0, 0, 2](见下面代码中的示例)。以下代码可以运行，但对于大型数组来说速度很慢。

import numpy as np


def duplicates(ar):
    """
    Args:
        ar (array_like): array

    Returns:
        list of int: int is pointing to first occurence of unique value
    """
    # duplicates array:
    dup = np.full(ar.shape[0], -1, dtype=int)
    for i in range(ar.shape[0]):
        if dup[i] != -1:
            # i is already found to be a
            continue
        else:
            dup[i] = i
        for j in range(i + 1, ar.shape[0]):
            if (ar[i] == ar[j]).all():
                dup[j] = i
    return dup


if __name__ == '__main__':
    n = 100
    # shortest extreme for n points
    a1 = np.array([[0, 1, 2]] * n)
    assert (duplicates(a1) == np.full(n, 0)).all(), True

    # longest extreme for n points
    a2 = np.linspace(0, 1, n * 3).reshape((n, 3))
    assert (duplicates(a2) == np.arange(0, n)).all(), True

    # test case
    a3 = np.array([[1, 0, 0], [1, 0, 0], [2, 3, 4]])
    assert (duplicates(a3) == [0, 0, 2]).all(), True

知道如何加速这个过程(例如避免第二个 for 循环)或替代实现吗？干杯

Answer 1

您正在做的事情要求您在每个可能的配对中将 N 行(每行长度为 M)与另一行进行比较。这意味着在没有重复项的情况下，它充其量可以扩展为 O(N^2 * M)。

更好的方法是散列每一行。如果散列所需的时间缩放为 O(M)，那么这应该缩放为 O(N * M)。你可以用字典来做到这一点:

def duplicates(ar):
    """
    Args:
        ar (array_like): array

    Returns:
        list of int: int is pointing to first occurence of unique value
    """
    first_occurence = {}
    # duplicates array:
    dup = np.zeros(ar.shape[0], dtype=int)
    for i in range(ar.shape[0]):
        as_tuple = tuple(ar[i])
        if as_tuple not in first_occurence:
            first_occurence[as_tuple] = i
        dup[i] = first_occurence[as_tuple]
    return dup