I am trying to solve the following problem efficiently:
Given the following class
class Obj:
t1: int
t2: int
def __hash__(self):
# compute hash on frozenset to make sure switching the to and from id has no effect
return hash(frozenset((self.t1, self.t2)))
def __eq__(self, other: 'Obj'):
pairs = (self.t1, self.t2)
return other.r1 in pairs and other.t2 in pairs
Hence, Obj(t1=1, t2=2) == Obj(t1=2, t2=1)
So given a set of Obj
as in s = {Obj(1, 2), Obj(1, 3), Obj(4, 1), Obj(9, 5)}
how can I find all objects which have either t1=1
OR t2=1
? So from from the set above I only would get [Obj(1, 2), Obj(1, 3), Obj(4, 1)]
What I have tried so far:
My current approach is to sort the set twice. Once with s_t1 = sorted(s, lambda x: x.t1)
and s_t2 = sorted(s, lambda x: x.t2)
.
I used sortedcontainers
for convenience but one could also use classical bisect
key = 1
from sortedcontainers import SortedKeyList
s_t1 = SortedKeyList(s, key=lambda x: x.t1)
s_t2 = SortedKeyList(s, key=lambda x: x.t2)
t1_results = s_t1[s_t1.bisect_key_left(key):s_t1.bisect_key_right(key)]
t2_results = s_t2[s_t2.bisect_key_left(key):s_t2.bisect_key_right(key)]
output = t1_restuls + t2_results
However, this feels inefficient since I need to sort twice.
Hence I am wondering what would be a better and faster (!) approach to do this. In particular, how can I make better use of the fact that Obj(t1=1, t2=2) == Obj(t1=2, t2=1)
when searching the list?