优化大集合的交集

Question

The premise is simple: I have two integers, a and b and I want to find i st a + i and b + i are both in a given list. The list rs is very large (10e9 items). I have the following code:

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

The issue at hand is that a1 and b1 are pre-computed first which creates a memory problem. Can I optimize my code somehow? General comments about the method are also welcome.

Example input:

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

Expected output:

getlist(3,8) = [1]

Answer 1

You can optimize the memory usage by skipping the creation of the second set (and intermediate lists):

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]

The time and space complexity of this solution is O(n) .

Answer 2

If rs is already a set , this would be faster:

def getlist(a, b):
    return [i - a for i in rs if i > a and b + (i - a) in rs]

If it is not, then you have to make the set first (otherwise the above algorithm would be very slow) and the performance is essentially the same as before:

def getlist(a, b):
    rs_set = set(rs)
    return [i - a for i in rs_set if i > a and b + (i - a) in rs_set]

However, if you are going to use the same function many times for different a and b values but the same rs , you can convert rs to a set once and reuse it every time.