python - 优化/删除循环

Question

我有以下一段代码，我想使用 numpy 对其进行优化，最好是删除循环。我看不出如何处理它，所以任何建议都会有所帮助。

indices 是一个 (N,2) numpy 整数数组，N 可以是几百万。代码所做的是在第一列中找到重复的索引。对于这些索引，我在第二列中对两个相应索引进行了所有组合。然后我将它们与第一列中的索引一起收集。

index_sets = []
uniques, counts = np.unique(indices[:,0], return_counts=True)
potentials = uniques[counts > 1]
for p in potentials:
    correspondents = indices[(indices[:,0] == p),1]
    combs = np.vstack(list(combinations(correspondents, 2)))
    combs = np.hstack((np.tile(p, (combs.shape[0], 1)), combs))
    index_sets.append(combs)

Answer 1

可以提出一些改进建议:

• 初始化输出数组，我们可以预先计算出存储每个组对应的组合所需的估计行数。我们知道对于 N 元素，可能的组合总数将为 N*(N-1)/2，从而为我们提供每个组的组合长度。此外，输出数组中的总行数将是所有这些间隔长度的总和。

• 在进入循环之前以矢量化方式预先计算尽可能多的东西。

• 使用循环获取组合，由于模式参差不齐，因此无法对其进行矢量化。使用 np.repeat 模拟平铺并在循环之前执行此操作，以便为我们提供每个组的第一个元素，从而为我们提供输出数组的第一列。

所以，考虑到所有这些改进，实现看起来像这样 -

# Remove rows with counts == 1 
_,idx, counts = np.unique(indices[:,0], return_index=True, return_counts=True)
indices = np.delete(indices,idx[counts==1],axis=0)

# Decide the starting indices of corresponding to start of new groups 
# charaterized by new elements along the sorted first column
start_idx = np.unique(indices[:,0], return_index=True)[1]
all_idx = np.append(start_idx,indices.shape[0])

# Get interval lengths that are required to store pairwise combinations
# of each group for unique ID from column-0
interval_lens = np.array([item*(item-1)/2 for item in np.diff(all_idx)])

# Setup output array and set the first column as a repeated array
out = np.zeros((interval_lens.sum(),3),dtype=int)
out[:,0] = np.repeat(indices[start_idx,0],interval_lens)

# Decide the start-stop indices for storing into output array 
ssidx = np.append(0,np.cumsum(interval_lens))

# Finally run a loop gto store all the combinations into initialized o/p array
for i in range(idx.size):
    out[ssidx[i]:ssidx[i+1],1:] = \
    np.vstack(combinations(indices[all_idx[i]:all_idx[i+1],1],2))

请注意，输出数组将是一个大的 (M, 3) 形状的数组，而不是像原始代码生成的那样拆分为数组列表。如果仍然需要这样，可以使用 np.split 来达到同样的目的。

此外，快速运行时测试表明所提议的代码没有太大改进。因此，可能大部分运行时间都花在了获取组合上。因此，似乎有替代方法 networkx特别适用于此类基于连接的问题可能更合适。