Python SortedList自定义对象查找技巧

本文探讨了在Python `sortedcontainers`库的`SortedList`中，如何高效且优雅地搜索自定义对象。针对直接使用字符串搜索自定义对象列表的挑战，文章提出了一种通过在自定义类中实现富比较方法（如`__lt__`）来处理与字符串的比较，从而使`bisect_left`等方法能够直接接受搜索字符串的解决方案。这种方法避免了创建临时对象或手动实现二分查找的繁琐。

在Python开发中，当我们需要维护一个大型的、有序的自定义对象集合时，sortedcontainers库提供的SortedList是一个非常强大的工具。它结合了列表的便利性和平衡二叉树的查找效率，使得插入、删除和查找操作都具有对数时间复杂度。然而，当我们需要根据自定义对象的某个特定属性（例如，一个名称字符串）来搜索列表时，常常会遇到一些挑战。

遇到的问题与常见误区

假设我们有一个Supplier类，包含Name、Id和SapId等属性，并且我们希望根据Name属性在SortedList中查找供应商。

from typing import List
from sortedcontainers import SortedList

class Supplier:
    def __init__(self, name: str, id: int, sap_id: int):
        self.Name = name
        self.Id = id
        self.SapId = sap_id

    def __repr__(self):
        # 便于调试和显示
        return f"Supplier(Name='{self.Name}', Id={self.Id})"

class Data:
    def __init__(self):
        # 初始化时指定key，按名称小写排序
        self.suppliers = SortedList(key=lambda x: x.Name.lower())

    def find_supplier(self, name: str):
        # 尝试直接使用bisect_left搜索字符串
        # 这里的name是str类型，而SortedList期望Supplier类型
        index = self.suppliers.bisect_left(name)
        if index != len(self.suppliers) and self.suppliers[index].Name.lower() == name.lower():
            return self.suppliers[index]
        return None

当我们尝试直接将一个字符串name传递给bisect_left方法时，会发现它无法正确工作。这是因为SortedList在内部比较元素时，即使指定了key函数，bisect_left仍然期望接收一个与列表元素类型兼容的对象进行比较，而不是key函数生成的比较值。换句话说，它会尝试将传入的str类型与列表中的Supplier类型进行比较，这通常会导致类型错误或不符合预期的行为。

一种常见的“变通”方法是创建一个临时的Supplier对象，只填充搜索所需的Name属性，然后用这个临时对象进行搜索：

    # Data类中的find_supplier方法（不推荐）
    def find_supplier_ugly(self, name: str):
        # 创建一个临时Supplier对象进行搜索
        temporary_supplier = Supplier(name, 0, 0) # Id和SapId可以是任意值
        index = self.suppliers.bisect_left(temporary_supplier)
        if index != len(self.suppliers) and self.suppliers[index].Name.lower() == name.lower():
            return self.suppliers[index]
        return None

虽然这种方法能够实现功能，但它显得不够优雅，且在每次搜索时都创建了一个无实际业务意义的临时对象，增加了代码的复杂性和潜在的性能开销。

优雅的解决方案：实现富比较方法

Python的面向对象特性允许我们通过实现“富比较方法”（rich comparison methods）来定义对象之间的比较行为。这些方法包括__lt__ (小于), __le__ (小于等于), __eq__ (等于), __ne__ (不等于), __gt__ (大于), __ge__ (大于等于)。通过在Supplier类中实现这些方法，我们可以让Supplier对象知道如何与另一个Supplier对象进行比较，甚至是如何与一个字符串进行比较。

当SortedList在没有key函数的情况下初始化时，它会依赖于其元素的自然比较顺序，即通过调用元素的富比较方法来确定排序。因此，我们可以利用这一点，让Supplier对象能够直接与字符串进行比较。

修改Supplier类

我们将修改Supplier类，使其能够与字符串进行比较。这里我们主要实现__lt__方法，因为bisect_left主要依赖于小于比较来确定插入位置。

class Supplier:
    def __init__(self, name: str, id: int = 0, sap_id: int = 0): # 默认值，方便临时对象创建（如果需要）
        self.Name = name
        self.Id = id
        self.SapId = sap_id

    def __repr__(self):
        return f"Supplier('{self.Name}')" # 更简洁的表示

    # 实现小于比较方法
    def __lt__(self, other):
        if isinstance(other, str):
            # 如果other是字符串，则将Supplier的Name与字符串比较
            return self.Name.lower() < other.lower()
        elif isinstance(other, Supplier):
            # 如果other是另一个Supplier对象，则比较它们的Name
            return self.Name.lower() < other.Name.lower()
        # 否则，抛出TypeError或返回NotImplemented，取决于具体需求
        return NotImplemented

    # 同样，为了完整性和健壮性，建议实现__eq__
    def __eq__(self, other):
        if isinstance(other, str):
            return self.Name.lower() == other.lower()
        elif isinstance(other, Supplier):
            return self.Name.lower() == other.Name.lower()
        return NotImplemented

修改Data类和搜索方法

在Supplier类定义了比较行为后，Data类初始化SortedList时就不再需要key参数了，因为SortedList会直接使用Supplier对象自身的比较逻辑。find_supplier方法现在可以直接将搜索字符串传递给bisect_left。

class Data:
    def __init__(self):
        # SortedList不再需要key参数，因为它会使用Supplier对象的__lt__方法
        self.suppliers = SortedList()

    def find_supplier(self, name: str):
        # bisect_left现在可以直接接收字符串，因为Supplier定义了与字符串的比较
        index = self.suppliers.bisect_left(name)

        # 检查找到的索引是否有效，并且元素名称是否完全匹配（考虑大小写）
        if index != len(self.suppliers) and self.suppliers[index].Name.lower() == name.lower():
            return self.suppliers[index]

        return None

完整示例与验证

下面是一个完整的示例，演示了如何使用这种方法：

from sortedcontainers import SortedList

class Supplier:
    def __init__(self, name: str, id: int = 0, sap_id: int = 0):
        self.Name = name
        self.Id = id
        self.SapId = sap_id

    def __repr__(self):
        return f"Supplier('{self.Name}')"

    def __lt__(self, other):
        if isinstance(other, str):
            return self.Name.lower() < other.lower()
        elif isinstance(other, Supplier):
            return self.Name.lower() < other.Name.lower()
        return NotImplemented

    def __eq__(self, other):
        if isinstance(other, str):
            return self.Name.lower() == other.lower()
        elif isinstance(other, Supplier):
            return self.Name.lower() == other.Name.lower()
        return NotImplemented

class Data:
    def __init__(self):
        self.suppliers = SortedList()

    def find_supplier(self, name: str):
        index = self.suppliers.bisect_left(name)
        if index != len(self.suppliers) and self.suppliers[index].Name.lower() == name.lower():
            return self.suppliers[index]
        return None

# 示例使用
d = Data()
d.suppliers.add(Supplier('Apple', 101, 1001))
d.suppliers.add(Supplier('Banana', 102, 1002))
d.suppliers.add(Supplier('apple', 103, 1003)) # 故意添加一个同名但ID不同的
d.suppliers.add(Supplier('Cherry', 104, 1004))

print("SortedList内容:", d.suppliers)

# 搜索存在的供应商
found_supplier_a = d.find_supplier('Apple')
print(f"搜索 'Apple': {found_supplier_a}") # 预期输出 Supplier('Apple')

found_supplier_b = d.find_supplier('banana')
print(f"搜索 'banana': {found_supplier_b}") # 预期输出 Supplier('Banana')

# 搜索不存在的供应商
found_supplier_d = d.find_supplier('Durian')
print(f"搜索 'Durian': {found_supplier_d}") # 预期输出 None

# 搜索与现有名称大小写不同的，但实际存在的
found_supplier_upper_a = d.find_supplier('APPLE')
print(f"搜索 'APPLE': {found_supplier_upper_a}") # 预期输出 Supplier('Apple')

输出结果：

SortedList内容: [Supplier('Apple'), Supplier('apple'), Supplier('Banana'), Supplier('Cherry')]
搜索 'Apple': Supplier('Apple')
搜索 'banana': Supplier('Banana')
搜索 'Durian': None
搜索 'APPLE': Supplier('Apple')

从输出可以看出，bisect_left成功地定位到了元素，并且find_supplier方法能够正确地返回或判断为None。

注意事项与总结

1. 比较逻辑一致性： 确保__lt__和__eq__中的比较逻辑与你期望的排序和查找逻辑一致（例如，都使用.lower()进行大小写不敏感比较）。
2. NotImplemented的正确使用： 当无法处理与other类型的比较时，返回NotImplemented是最佳实践。这允许Python尝试other对象的反向操作（例如，str.__gt__(self)），如果仍然无法处理，才会抛出TypeError。
3. SortedList初始化： 采用此方法后，SortedList在初始化时不再需要key参数，因为它会依赖于元素自身的比较方法。
4. 最终相等性检查： 即使bisect_left找到了一个可能的索引，最后一步的self.suppliers[index].Name.lower() == name.lower()检查仍然是至关重要的。这是因为bisect_left只保证找到一个“插入点”，在这个点之前的所有元素都小于或等于搜索值，而在这个点之后的所有元素都大于搜索值。它并不保证找到的元素就是精确匹配的，尤其是在存在重复元素或部分匹配的情况下。
5. 灵活性： 这种方法不仅适用于字符串，你也可以在富比较方法中处理其他自定义类型，从而实现更复杂的搜索逻辑。

通过在自定义类中实现富比较方法，我们能够以一种更Pythonic、更优雅的方式解决SortedList中自定义对象的搜索问题，避免了不必要的临时对象创建，并使代码更加清晰和易于维护。