Python中从集合获取匹配自定义对象的方法（重写__eq__/__hash__）

我需要检查存储自定义对象的集合中是否存在哈希匹配的对象，若存在则获取该对象的引用。以下是我的初始代码实现：

class SetObject():
    def __init__(
        self,
        a: int,
        b: int,
        c: int
    ):
        self.a = a
        self.b = b
        self.c = c

    def __repr__(self):
        return f"SetObject ({self.a} {self.b} {self.c}) (id: {id(self)})"

    def __eq__(self, other):
        return isinstance(other, SetObject) \
               and self.__hash__() == other.__hash__()

    def __hash__(self):
        # Hash only depends on a, b
        return hash(
            (self.a,self.b)
        )

x = SetObject(1,2,3)
y = SetObject(4,5,6)

object_set = set([x,y])

print(f"{object_set=}")

z = SetObject(1,2,7)
print(f"{z=}")
if z in object_set:
    print("Is in set")
    # Get the object in set which is equal to z
    for element in object_set:
       if element == z:
           print(element)
           z = element

print(f"{z=}")

问题分析 #

1. __eq__实现缺陷：直接通过哈希值判断相等存在风险——不同对象可能出现哈希碰撞，导致错误判定相等。正确的做法应该直接比较决定哈希的核心字段（这里是a和b）。
2. 遍历集合效率低：原代码通过遍历整个集合查找匹配对象，时间复杂度为O(n)，没有利用集合O(1)查找的特性优势。

优化方案 #

1. 修正__eq__并使用生成器快速获取对象

先修正__eq__方法，确保相等性判断准确：

def __eq__(self, other):
    return isinstance(other, SetObject) and self.a == other.a and self.b == other.b

然后利用生成器表达式，找到第一个匹配对象后立即停止遍历：

if z in object_set:
    print("Is in set")
    # 快速获取集合中匹配的对象
    z = next(element for element in object_set if element == z)

2. 用字典实现O(1)查找（推荐频繁查找场景）

如果需要多次执行查找操作，建议维护一个辅助字典，以(a,b)元组为键，对应集合中的对象为值：

# 初始化集合与字典
x = SetObject(1,2,3)
y = SetObject(4,5,6)
object_set = {x, y}
object_map = {(obj.a, obj.b): obj for obj in object_set}

# 查找操作
z = SetObject(1,2,7)
print(f"{z=}")
if (z.a, z.b) in object_map:
    print("Is in set")
    z = object_map[(z.a, z.b)]

这种方式的查找时间复杂度为O(1)，是效率最高的方案。

内容的提问来源于stack exchange，提问作者Alessi 42