如何获取NumPy数组中重复元素的所有索引列表

Question

我正在尝试获取numpy数组中所有重复元素的索引，但我目前找到的解决方案对于大型（>20000个元素）输入数组非常低效（需要大约9秒钟）。这里的思路很简单：

1. records_array 是一个包含时间戳（datetime）的numpy数组，我们想要提取出其中重复时间戳的索引。

2. time_array 是一个numpy数组，包含在records_array中重复的所有时间戳。

3. records 是一个django QuerySet（可以轻松转换为列表），其中包含一些Record对象。我们想要创建一个由所有可能的Record的tagId属性组成的二元组列表，这些属性对应于从records_array中发现的重复时间戳。

这是我目前使用的有效代码，但效率不高：

tag_couples = [];
for t in time_array:
    users_inter = np.nonzero(records_array == t)[0] # Get all repeated timestamps in records_array for time t
    l = [str(records[i].tagId) for i in users_inter] # Create a temporary list containing all tagIds recorded at time t
    if l.count(l[0]) != len(l): #remove tuples formed by the first tag repeated
        tag_couples +=[x for x in itertools.combinations(list(set(l)),2)] # Remove duplicates with list(set(l)) and append all possible couple combinations to tag_couples

我相信可以通过使用Numpy进行优化，但是我找不到一种方法来比较records_array和time_array中的每个元素而不使用for循环（不能只使用==比较，因为它们都是数组）。

- morens

---

关于所问的内容：如果你只是想要删除重复项，pandas.DataFrame.drop_duplicates() 可以帮你省去自己查找重复项的麻烦。 - undefined

Answer 1

使用numpy的向量化解决方案，利用unique()函数的魔力。

import numpy as np

# create a test array
records_array = np.array([1, 2, 3, 1, 1, 3, 4, 3, 2])

# creates an array of indices, sorted by unique element
idx_sort = np.argsort(records_array)

# sorts records array so all unique elements are together 
sorted_records_array = records_array[idx_sort]

# returns the unique values, the index of the first occurrence of a value, and the count for each element
vals, idx_start, count = np.unique(sorted_records_array, return_counts=True, return_index=True)

# splits the indices into separate arrays
res = np.split(idx_sort, idx_start[1:])

#filter them with respect to their size, keeping only items occurring more than once
vals = vals[count > 1]
res = filter(lambda x: x.size > 1, res)

以下是原始答案的代码，需要使用更多的内存，使用numpy广播并调用unique两次：

---

以下是原始答案的代码，需要使用更多的内存，使用numpy广播并调用unique两次：

records_array = array([1, 2, 3, 1, 1, 3, 4, 3, 2])
vals, inverse, count = unique(records_array, return_inverse=True,
                              return_counts=True)

idx_vals_repeated = where(count > 1)[0]
vals_repeated = vals[idx_vals_repeated]

rows, cols = where(inverse == idx_vals_repeated[:, newaxis])
_, inverse_rows = unique(rows, return_index=True)
res = split(cols, inverse_rows[1:])

如预期的那样，res = [array([0, 3, 4]), array([1, 8]), array([2, 5, 7])]

- gg349

---

作为此解决方案的警告，对于具有重复项的数组，idx_sort 是不确定性的，因为 np.sort 中的默认排序是不稳定的。要解决这个问题，请确保使用合并排序类型：np.argsort(a, kind='mergesort')。 - user27443

Answer 2

• 答案比较复杂，取决于数组中唯一元素的数量和大小。
• 以下测试：
  • 测试包含2M个元素，最多20k个唯一元素的数组。
  • 测试包含不超过80k个元素，最大20k个唯一元素的数组。
    • 对于小于40k元素的数组，测试中唯一元素的数量最多为数组大小的一半（例如10k元素的数组将具有最多5k个唯一元素）。

包含2M个元素的数组

• 对于不超过约200个唯一元素的情况，np.where比defaultdict更快，但比pandas.core.groupby.GroupBy.indices和np.unique慢。
• 使用pandas的解决方案是处理大型数组的最快解决方案。

包含不超过80k个元素的数组

• 这更取决于数组的大小和唯一元素的数量。
• defaultdict是处理不超过2400个元素的数组的快速选项，特别是当唯一元素很多时。
• 对于包含40k个以上元素和20k个唯一元素的数组，pandas是最快的选项。

%timeit

import random
import numpy
import pandas as pd
from collections import defaultdict

def dd(l):
    # default_dict test
    indices = defaultdict(list)
    for i, v in enumerate(l):
        indices[v].append(i)
    return indices


def npw(l):
    # np_where test
    return {v: np.where(l == v)[0] for v in np.unique(l)}


def uni(records_array):
    # np_unique test
    idx_sort = np.argsort(records_array)
    sorted_records_array = records_array[idx_sort]
    vals, idx_start, count = np.unique(sorted_records_array, return_counts=True, return_index=True)
    res = np.split(idx_sort, idx_start[1:])
    return dict(zip(vals, res))


def daf(l):
    # pandas test
    return pd.DataFrame(l).groupby([0]).indices


data = defaultdict(list)

for x in range(4, 20000, 100):  # number of unique elements
    # create 2M element list
    random.seed(365)
    a = np.array([random.choice(range(x)) for _ in range(2000000)])
    
    res1 = %timeit -r2 -n1 -q -o dd(a)
    res2 = %timeit -r2 -n1 -q -o npw(a)
    res3 = %timeit -r2 -n1 -q -o uni(a)
    res4 = %timeit -r2 -n1 -q -o daf(a)
    
    data['defaut_dict'].append(res1.average)
    data['np_where'].append(res2.average)
    data['np_unique'].append(res3.average)
    data['pandas'].append(res4.average)
    data['idx'].append(x)

df = pd.DataFrame(data)
df.set_index('idx', inplace=True)

df.plot(figsize=(12, 5), xlabel='unique samples', ylabel='average time (s)', title='%timeit test: 2 run 1 loop each')
plt.legend(bbox_to_anchor=(1.0, 1), loc='upper left')
plt.show()

测试200万个元素的结果

测试高达80k个元素的结果

- Trenton McKinney

Answer 3

您也可以这样做：

a = [1,2,3,1,1,3,4,3,2]
index_sets = [np.argwhere(i==a) for i in np.unique(a)]

这将为您提供一组带有唯一元素索引的数组。

[array([[0],[3],[4]], dtype=int64), 
array([[1],[8]], dtype=int64), 
array([[2],[5],[7]], dtype=int64), 
array([[6]], dtype=int64)]

新增：进一步改变列表推导式还可以丢弃单个独特值，并在存在许多唯一单次发生元素的情况下解决速度问题：

new_index_sets = [np.argwhere(i[0]== a) for i in np.array(np.unique(a, return_counts=True)).T if i[1]>=2]

这将会产生：

[array([[0],[3],[4]], dtype=int64), 
 array([[1],[8]], dtype=int64), 
 array([[2],[5],[7]], dtype=int64)]

- Ashish

---

如果数组包含许多唯一值，这将非常缓慢。 - gg349

@ gg349 感谢指出。一个小修改也解决了这个问题。现在我希望速度足够快。 - Ashish

代码现在变慢了，因为它调用where()的次数与重复值的数量一样多。每次调用where都必须遍历整个数组。 - gg349

在我的电脑上，使用输入random_integers(-1000,1000,20000)，我使用上述第一种解决方案大约可以获得16倍的加速。我猜你坚持使用了一个只有很少不同值的随机数组？ - gg349

确实。我刚刚使用了你的record_array并将其乘以2000，即records_array = np.array([1, 2, 3, 1, 1, 3, 4, 3, 2, 1]*2000)，导致了这种情况。使用np.random.randint(-1000,1000,20000)，你的解决方案快了约16倍。 - Ashish

Answer 4

我发现不使用np.unique，而是使用np.diff显著更快，并可以更好地处理非排序初始数组。

为了证明这一点，我运行了@Trenton McKinney的基准测试来演示差分解决方案超越其他方法。它也不需要一个已排序的数组或对数组进行排序，这是一个显著的优势。

以下是这个函数：

def find_repeats(arr: np.ndarray) -> np.ndarray:
    """Find indices of repeat values in an array.

    Args:
        arr (np.ndarray): An array to find repeat values in.

    Returns:
        np.ndarray: An array of indices into arr which are the values which
            repeat.
    """

    arr_diff = np.diff(arr, append=[arr[-1] + 1])
    res_mask = arr_diff == 0
    arr_diff_zero_right = np.nonzero(res_mask)[0] + 1
    res_mask[arr_diff_zero_right] = True
    return np.nonzero(res_mask)[0]

200万个元素

20k个元素

完整测试代码

import random
from matplotlib import pyplot as plt
import numpy as np
import pandas as pd
from collections import defaultdict
import time


def find_repeats(arr: np.ndarray) -> np.ndarray:
    """Find indices of repeat values in an array.

    Args:
        arr (np.ndarray): An array to find repeat values in.

    Returns:
        np.ndarray: An array of indices into arr which are the values which
            repeat.
    """

    arr_diff = np.diff(arr, append=[arr[-1] + 1])
    res_mask = arr_diff == 0
    arr_diff_zero_right = np.nonzero(res_mask)[0] + 1
    res_mask[arr_diff_zero_right] = True
    return np.nonzero(res_mask)[0]


def dd(l):
    # default_dict test
    indices = defaultdict(list)
    for i, v in enumerate(l):
        indices[v].append(i)
    return indices


def npw(l):
    # np_where test
    return {v: np.where(l == v)[0] for v in np.unique(l)}


def uni(records_array):
    # np_unique test
    idx_sort = np.argsort(records_array)
    sorted_records_array = records_array[idx_sort]
    vals, idx_start, count = np.unique(
        sorted_records_array, return_counts=True, return_index=True)
    res = np.split(idx_sort, idx_start[1:])
    return dict(zip(vals, res))


def daf(l):
    # pandas test
    return pd.DataFrame(l).groupby([0]).indices


data = defaultdict(list)

for x in range(4, 20000, 1000):  # number of unique elements
    print(f"{x} trial done")
    # create 2M element list
    random.seed(365)
    a = np.array([random.choice(range(x)) for _ in range(2000000)])
    num_runs = 2
    t0 = time.time()
    for i in range(num_runs):
        dd(a)
    res1 = time.time() - t0

    t0 = time.time()
    for i in range(num_runs):
        uni(a)
    res3 = time.time() - t0

    t0 = time.time()
    for i in range(num_runs):
        daf(a)
    res4 = time.time() - t0

    t0 = time.time()
    for i in range(num_runs):
        find_repeats(a)
    res5 = time.time() - t0

    data['defaut_dict'].append(res1 / num_runs)
    data['np_unique'].append(res3 / num_runs)
    data['pandas'].append(res4 / num_runs)
    data['np_diff'].append(res5 / num_runs)
    data['idx'].append(x)

df = pd.DataFrame(data)
df.set_index('idx', inplace=True)

df.plot(figsize=(12, 5), xlabel='unique samples',
        ylabel='average time (s)', title='%timeit test: 2 run 1 loop each')
plt.legend(bbox_to_anchor=(1.0, 1), loc='upper left')
plt.show()

- user27443

---

这个基准测试有缺陷，因为np.diff总是返回一个空数组。它是唯一一个返回错误答案的方法。请进行更新。 - undefined

Answer 5

你可以这样做：

您可以沿以下方向进行操作：

1. add original index ref so [[1,0],[2,1],[3,2],[1,3],[1,4]...
2. sort on [:,0]
3. use np.where(ra[1:,0] != ra[:-1,0])
4. use the list of indexes from above to construct your final list of lists

编辑 - 好的，我在快速回复后离开了一段时间，现在看到我的答案被投票下降了，这是公平的，因为numpy.argsort()比我的建议要好得多。我确实投票支持了numpy.unique()的答案，因为这是一个有趣的功能。然而，如果你使用timeit，你会发现

idx_start = np.where(sorted_records_array[:-1] != sorted_records_array[1:])[0] + 1
res = np.split(idx_sort, idx_start)

比起其他，略微快一些

vals, idx_start = np.unique(sorted_records_array, return_index=True)
res = np.split(idx_sort, idx_start[1:])

以下是对@Nicolas提出问题的进一步编辑

我不确定你能否做到。可以获取两个与断点相对应的索引数组，但无法使用np.split将数组的不同“行”分成不同大小的块。

a = np.array([[4,27,42,12, 4 .. 240, 12], [3,65,23...] etc])
idx = np.argsort(a, axis=1)
sorted_a = np.diagonal(a[:, idx[:]]).T
idx_start = np.where(sorted_a[:,:-1] != sorted_a[:,1:])

# idx_start => (array([0,0,0,..1,1,..]), array([1,4,6,7..99,0,4,5]))

但是根据您想要使用信息的目的，这可能已经足够好了。

- paddyg

---

如果原始数组具有多行，并且在没有for循环的情况下迭代遍历每一行，是否可以通过某种方式使其工作？ - Nickpick

你无法进行拆分。我会在上面添加一个编辑。 - paddyg

Answer 6

所以我无法摆脱for循环，但是我能够使用set数据类型和list.count()方法将其缩小到对for循环的轻微使用：

data = [1,2,3,1,4,5,2,2]
indivs = set(data)

multi_index = lambda lst, val: [i for i, x in enumerate(lst) if x == val]

if data != list(indivs):
    dupes = [multi_index(data, i) for i in indivs if data.count(i) > 1]

在循环遍历indivs集合时，该集合包含值（无重复项），然后在完整列表中循环遍历，如果找到重复项，则执行操作。如果这种方法不够快，可以考虑使用numpy替代方案。如果需要加速，生成器对象也可能有所帮助。

编辑：gg349的答案提供了我正在研究的numpy解决方案！

- Matthew

Answer 7

np.unique 适用于所有指数

@gg349的解决方案封装成一个函数:

def np_unique_indices(arr, **kwargs):
    """Unique indices for N-D arrays."""
    vals, indices, *others = np_unique_indices_1d(arr.reshape(-1), **kwargs)
    indices = [np.stack(np.unravel_index(x, arr.shape)) for x in indices]
    return vals, indices, *others


def np_unique_indices_1d(arr, **kwargs):
    """Unique indices for 1D arrays."""
    sort_indices = np.argsort(arr)
    arr = np.asarray(arr)[sort_indices]
    vals, first_indices, *others = np.unique(
        arr, return_index=True, **kwargs
    )
    indices = np.split(sort_indices, first_indices[1:])
    for x in indices:
        x.sort()
    return vals, indices, *others

它本质上与np.unique相同，但返回所有索引，而不仅仅是第一个索引。

---

使用示例：

arr = np.array([
    [0, 1, 1, 0],
    [0, 2, 2, 0],
    [0, 2, 2, 0],
    [0, 1, 1, 0],
])

vals, indices = np_unique_indices(arr)

for val, idx in zip(vals, indices):
    print(f"{val}:\n{idx}\n")

输出：

0:
[[0 0 1 1 2 2 3 3]
 [0 3 0 3 0 3 0 3]]

1:
[[0 0 3 3]
 [1 2 1 2]]

2:
[[1 1 2 2]
 [1 2 1 2]]

- Mateen Ulhaq

Answer 8

import numpy as np
from numpy.lib import recfunctions as rfn

ndtype = [('records_array', int)] # Check the data type
records_array = np.ma.array([1, 2, 1, 3, 2, 3, 3, 4, 5]).view(ndtype) # Structured array
idxs = list(rfn.find_duplicates(records_array, key=None, ignoremask=True, return_index=True)[1]) # List of indices of repeated elements

- Leonardo Queiroz

Answer 9

numba.jit

另一种解决方案，但使用 numba.jit：

def np_unique_indices(arr, **kwargs):
    """Unique indices for N-D arrays."""
    vals, indices = np_unique_indices_1d(arr.reshape(-1))
    vals = np.asarray(vals)
    indices = [np.stack(np.unravel_index(x, arr.shape)) for x in indices]
    return vals, indices


@numba.njit
def np_unique_indices_1d(arr):
    """Unique indices for 1D arrays."""
    idxs = [[0 for _ in range(0)] for _ in range(0)]
    ptr = {}
    ptr_count = 0

    for i, x in enumerate(arr):
        if (x in ptr) == False:
            idxs.append([0 for _ in range(0)])
            ptr[x] = ptr_count
            ptr_count += 1
        idxs[ptr[x]].append(i)

    vals = [x for x in ptr]
    idxs = [np.array(x) for x in idxs]
    return vals, idxs

使用 @Trenton McKinney 和 user27443 的基准测试：

请注意，所有解决方案的性能都取决于数组的大小和唯一标签的数量，因此我建议您根据自己的数据进行测试。

- Mateen Ulhaq