python - 查找多个数组具有相同值的索引的快速算法

Question

我正在寻找加速(或替换)我的数据分组算法的方法。

我有一个 numpy 数组列表。我想生成一个新的 numpy 数组，这样这个数组的每个元素对于每个索引都是相同的，而原始数组也是相同的。 (如果不是这种情况则有所不同。)

这听起来有点尴尬，所以举个例子:

# Test values:
values = [
    np.array([10, 11, 10, 11, 10, 11, 10]),
    np.array([21, 21, 22, 22, 21, 22, 23]),
    ]

# Expected outcome: np.array([0, 1, 2, 3, 0, 3, 4])
#                             *           *

请注意，我标记的预期结果的元素(索引 0 和 4)具有相同的值(0)，因为原来的两个数组也相同(即 10 和 21)。索引为 3 和 5 (3) 的元素类似。

该算法必须处理任意数量(大小相等)的输入数组，并且还为每个结果数字返回它们对应的原始数组的值。 (所以对于这个例子，“3”指的是 (11, 22)。)

这是我目前的算法:

import numpy as np

def groupify(values):
    group = np.zeros((len(values[0]),), dtype=np.int64) - 1 # Magic number: -1 means ungrouped.
    group_meanings = {}
    next_hash = 0
    matching = np.ones((len(values[0]),), dtype=bool)
    while any(group == -1):
        this_combo = {}

        matching[:] = (group == -1)
        first_ungrouped_idx = np.where(matching)[0][0]

        for curr_id, value_array in enumerate(values):
            needed_value = value_array[first_ungrouped_idx]
            matching[matching] = value_array[matching] == needed_value
            this_combo[curr_id] = needed_value
        # Assign all of the found elements to a new group
        group[matching] = next_hash
        group_meanings[next_hash] = this_combo
        next_hash += 1
    return group, group_meanings

请注意，表达式 value_array[matching] == needed_value 会针对每个单独的索引多次求值，这就是缓慢的来源。

我不确定我的算法是否可以进一步加速，但我也不确定它是否是开始时的最佳算法。有更好的方法吗？

Answer 1

终于破解了矢量化解决方案!这是一个有趣的问题。问题是我们必须标记从列表的相应数组元素中获取的每一对值。然后，我们应该根据它们在其他对中的唯一性来标记每个这样的对。因此，我们可以使用 np.unique 滥用其所有可选参数，最后做一些额外的工作来保持最终输出的顺序。下面是基本上分三个阶段完成的实现-

# Stack as a 2D array with each pair from values as a column each.
# Convert to linear index equivalent considering each column as indexing tuple
arr = np.vstack(values)
idx = np.ravel_multi_index(arr,arr.max(1)+1)

# Do the heavy work with np.unique to give us :
#   1. Starting indices of unique elems, 
#   2. Srray that has unique IDs for each element in idx, and 
#   3. Group ID counts
_,unq_start_idx,unqID,count = np.unique(idx,return_index=True, \
                                        return_inverse=True,return_counts=True)

# Best part happens here : Use mask to ignore the repeated elems and re-tag 
# each unqID using argsort() of masked elements from idx
mask = ~np.in1d(unqID,np.where(count>1)[0])
mask[unq_start_idx] = 1
out = idx[mask].argsort()[unqID]

运行时测试

让我们将建议的矢量化方法与原始代码进行比较。由于提议的代码只为我们提供组 ID，因此为了公平的基准测试，我们只删除原始代码中未用于提供给我们的部分。所以，这里是函数定义 -

def groupify(values):  # Original code
    group = np.zeros((len(values[0]),), dtype=np.int64) - 1
    next_hash = 0
    matching = np.ones((len(values[0]),), dtype=bool)
    while any(group == -1):

        matching[:] = (group == -1)
        first_ungrouped_idx = np.where(matching)[0][0]

        for curr_id, value_array in enumerate(values):
            needed_value = value_array[first_ungrouped_idx]
            matching[matching] = value_array[matching] == needed_value
        # Assign all of the found elements to a new group
        group[matching] = next_hash
        next_hash += 1
    return group

def groupify_vectorized(values):  # Proposed code
    arr = np.vstack(values)
    idx = np.ravel_multi_index(arr,arr.max(1)+1)
    _,unq_start_idx,unqID,count = np.unique(idx,return_index=True, \
                                        return_inverse=True,return_counts=True)    
    mask = ~np.in1d(unqID,np.where(count>1)[0])
    mask[unq_start_idx] = 1
    return idx[mask].argsort()[unqID]

具有大型数组的列表上的运行时结果 -

In [345]: # Input list with random elements
     ...: values = [item for item in np.random.randint(10,40,(10,10000))]

In [346]: np.allclose(groupify(values),groupify_vectorized(values))
Out[346]: True

In [347]: %timeit groupify(values)
1 loops, best of 3: 4.02 s per loop

In [348]: %timeit groupify_vectorized(values)
100 loops, best of 3: 3.74 ms per loop