元组列表：仅根据第一个或第二个位置的公共元素选择元组

Question

我有一个巨大的元组列表：

ijs = [(0,1),(0,2),(0,3), (3,2)...]

对于给定值v，我只想获取具有i=v 或j=v 的(i,j) 对（来自存储在ijs 中的所有可能的(i,j) 对）。

例如对于v=0 并给出ijs = [(0,1),(0,2),(0,3), (3,2)]，那么我应该返回only_current = [(0,1),(0,2),(0,3)]

---

示例：

请忽略前 3 行，我在其中构建了一个列表 ijs，其中包含元组。

import numpy as np

# IGNORE THIS PART UNTIL THE MAIN LOOP 
N= 1000
indices_upper_triangle = np.triu_indices(N, k = 1) # k=1 above diagonal
i,j = indices_upper_triangle
ijs = [(i,j) for i,j in zip(i,j)] # create (i,j) positions

# MAIN LOOP HERE 
# Main loop
all_neig_results = list()
for v in range(N): # for each point

    # from all possible (i,j) pairs, get only (i,j) pairs that have either i=v or j=v
    only_current = [item for item in ijs if v in item]

    all_neig_results.append(only_current)

循环中的列表理解超级慢。

%timeit [item for item in ijs if v in item]
15.9 s ± 361 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

如果我删除检查参数if v in item:

%timeit [item for item in ijs]
1.28 s ± 90.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

如何优化[item for item in ijs if v in item]？

Answer 1

你在这里做 O(N^2) 的工作，因为你每次都在内循环中遍历整个列表。

您应该预先计算一次：

import collections
def precompute(pairs):
    d = collections.defaultdict(list)
    for ij in pairs:
        d[ij[0]].append(ij)
        d[ij[1]].append(ij)
    return d

d = precompute(ijs)

使内部循环中的工作快速：

for v in range(N):
    only_current = d[v]

这样你只做 O(N) 的工作。

Answer 2

如果您的目标是找到最近的邻居，我建议您使用KDTree，它

可用于快速查找任意点的最近邻。

并避免遍历所有点/索引对。

见这个例子：How to find set of points in x,y grid using KDTree.query_ball_tree

关于leafsize 参数： KDTrees 将空间“分割”成不同的区域。每个叶子对应一个区域。如果leafsize =1，那么它会继续进行分割，直到每个区域只包含一个点。这意味着构建树需要很长时间，但搜索会很快。如果叶子包含更多点，则树更浅，因此构建它会更快。但是查询可能会更慢。文档中的“蛮力”位是指这样一个事实，即如果您在每个区域中有超过 1 个点，则需要蛮力找出该区域中最近的邻居。

leafsize=10 表示叶子有 10 个或更少。每次拆分都使得子节点具有（大约）父节点中一半的点，因此每个叶子中的精确点数将是可变的。还有其他实现细节可能会影响到这一点。