Python 中 in 操作符在集合与列表中的不同行为解析

本文深入探讨了 Python 中 in 操作符在列表和集合这两种数据结构中的不同行为。通过分析其内部实现机制，揭示了为何在某些情况下，使用列表会引发错误，而使用集合却能正常运行。此外，本文还提供了一个自定义类示例，用于更直观地理解 in 操作符的工作原理，并针对 PyTorch 张量比较问题，提出了相应的解决方案。

在 Python 中，in 操作符用于检查某个元素是否存在于一个集合（collection）中。然而，其行为会根据集合的类型（如列表、元组、集合、字典等）而有所不同。理解这些差异对于编写高效且无错的代码至关重要。

in 操作符的工作原理

x in collection 的行为取决于 collection 的类型。

1. 非哈希表集合（列表、元组等）

对于列表和元组等不使用哈希表的集合，in 操作符会遍历集合中的每个元素，并逐一进行比较。其内部逻辑类似于以下伪代码：

def is_in(x, collection):
  for c in collection:
      if (x is c or x == c):
          return True
  return False

首先比较 x 和 c 的身份（使用 is），如果身份相同，则返回 True。否则，比较它们的值是否相等（使用 ==）。这个过程会持续到找到第一个匹配项或遍历完整个集合。

立即学习“Python免费学习笔记（深入）”；

2. 哈希表集合（集合、字典等）

对于集合和字典等使用哈希表的集合，in 操作符的查找过程更加高效。其内部逻辑类似于以下伪代码：

def is_in(x, collection):
  # 选择集合中哈希值与 x 相同的元素子集
  subset = get_subset_by_hash(collection, hash(x))
  for c in subset:
      if (x is c or x == c):
          return True
  return False

首先，根据 x 的哈希值，从集合中选择一个子集，该子集包含所有哈希值与 x 相同的元素。然后，遍历这个子集，并进行身份和相等性比较。由于哈希表能够快速定位到可能的匹配项，因此查找速度通常比列表快得多。

注意： 集合中元素的哈希值是在元素添加到集合时计算的，而 x 的哈希值是在使用 in 操作符时计算的。

实例演示

为了更直观地理解 in 操作符的工作原理，我们可以创建一个自定义类 MyObj，并自定义其哈希计算逻辑和相等性比较逻辑：

class MyObj:
    def __init__(self, val, hashval):
        self._val = val
        self._hashval = hashval

    def __hash__(self):
        print(f"{str(self)} calling __hash__")
        return self._hashval

    def __eq__(self, other):
        print(f"{str(self)} calling __eq__, {other=}")
        return super().__eq__(other)

    def __repr__(self):
        return f"<{self.__class__.__name__}: {self._val}>"

然后，创建一些 MyObj 实例，并分别添加到集合和列表中：

MakeSong

AI音乐生成，生成高质量音乐，仅需30秒的时间

下载

a = MyObj("a", 123)
b = MyObj("b", 456)
d = MyObj("d", 456)  # 与 b 具有相同的哈希值

print("Creating set `s`")
s = set([a, b, d])

print("Creating list `lst`")
lst = [a, b, d]

通过观察输出，我们可以看到：

• 在创建集合时，Python 会计算每个元素的哈希值。
• 如果存在哈希冲突（例如，b 和 d 具有相同的哈希值），Python 还会调用 __eq__ 方法进行相等性比较。

接下来，我们可以使用 in 操作符来检查元素是否存在于集合和列表中：

>>> s
{, , }
>>> b in s
 calling __hash__
True
>>> d in s
 calling __hash__
 calling __eq__, other=
 calling __eq__, other=
True

>>> lst
[, , ]
>>> a in lst
True
>>> b in lst
 calling __eq__, other=
 calling __eq__, other=
True
>>> d in lst
 calling __eq__, other=
 calling __eq__, other=
 calling __eq__, other=
 calling __eq__, other=
True

通过观察输出，我们可以看到：

• 对于集合，Python 首先计算 x 的哈希值，然后查找哈希值相同的元素。如果存在哈希冲突，还会调用 __eq__ 方法进行相等性比较。
• 对于列表，Python 逐个比较元素，直到找到匹配项或遍历完整个列表。

PyTorch 张量比较问题

在 PyTorch 中，如果尝试比较大小不同的张量，会引发 RuntimeError。这是因为 PyTorch 会检查张量的形状是否兼容。

import torch
a = torch.Tensor(2,3)
b = torch.Tensor(2)

# case 1a:
# b  in list([a,a,b]) # raises an error: 
# RuntimeError: The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 0

# case 1b
b in set([a,a,b]) # True (i.e. no error)

当使用 in 操作符在列表中查找张量时，会按照顺序比较每个张量。如果遇到大小不同的张量，就会引发 RuntimeError。

而当使用 in 操作符在集合中查找张量时，由于集合使用哈希表，因此会首先比较哈希值。在 PyTorch 中，张量的哈希值实际上是其内存地址（id(self)）。因此，只有当 b 与集合中某个张量的内存地址相同时，才会进行相等性比较。由于 a 和 b 的内存地址不同，因此不会进行相等性比较，从而避免了 RuntimeError。

解决方案

为了解决这个问题，可以使用 torch.Tensor.size 属性，该属性是一个元组，表示张量的形状。可以创建一个字典，将不同形状的张量分别存储在不同的集合或列表中：

tensor_dict = {}
for tensor in [a, b]:
    size = tuple(tensor.size())
    if size not in tensor_dict:
        tensor_dict[size] = set()  # 或 list()
    tensor_dict[size].add(tensor)

# 检查 b 是否存在于具有相同形状的张量集合中
size_b = tuple(b.size())
if size_b in tensor_dict and b in tensor_dict[size_b]:
    print("b is in the collection")
else:
    print("b is not in the collection")

这种方法可以避免比较大小不同的张量，从而防止 RuntimeError。

总结

in 操作符在 Python 中的行为取决于集合的类型。理解其内部实现机制对于编写高效且无错的代码至关重要。对于不使用哈希表的集合（如列表和元组），in 操作符会逐个比较元素。对于使用哈希表的集合（如集合和字典），in 操作符会首先比较哈希值，然后再进行相等性比较。在处理 PyTorch 张量时，需要注意大小不同的张量比较可能引发 RuntimeError。可以使用 torch.Tensor.size 属性来避免这种问题。