python - 优化大集合的交集

Question

前提很简单:我有两个整数，a 和 b 我想找到 i s.t. a + i 和 b + i 都在给定列表中。列表 rs 非常大(10e9 项)。我有以下代码:

def getlist(a,b):
    a1 = set([i - a for i in rs if i>a])
    b1 = set([i-b for i in rs if i>b]) 

    tomp = list(a1.intersection(b1))
    return tomp

手头的问题是 a1 和 b1 首先被预先计算，这会产生内存问题。我能以某种方式优化我的代码吗？也欢迎对该方法提出一般性意见。

示例输入:

rs = [4,9,16]
a = 3
b = 8

预期输出:

getlist(3,8) = [1]

Answer 1

您可以通过跳过第二组(和中间列表)的创建来优化内存使用:

def getlist(a, b):
    a1 = {i - a for i in rs if i > a}
    return [i - b for i in rs if i > b and i - b in a1]

这个解决方案的时间和空间复杂度是O(n)。