Pandas多级索引下loc与dict查询的时间复杂度对比

实现方案 #

方案1：使用loc查询 #

import pandas

# 构造测试数据
df = pandas.DataFrame({"A": ["a", "a", "b", "b"],
                       "B": ["x", "y", "x", "y"],
                       "val": [1, 2, 3, 4]})
# 设置A、B列为两级索引
df = df.set_index(["A", "B"])

# 提取目标值
df.loc[("a", "x"), "val"]
# 输出：1

方案2：使用嵌套dict查询 #

# 构造测试数据
dat = {"a": {"x": 1, "y": 2},
       "b": {"x": 3, "y": 4}}

# 提取目标值
dat["a"]["x"]
# 输出：1

时间复杂度差异 #

• 嵌套dict的查询时间复杂度稳定为O(1)：Python的dict底层是哈希表实现，单次key查找都是常数时间，两次嵌套查找整体还是常数时间，且原生实现几乎没有额外开销，实际运行速度极快。
• 多级索引下的loc查询在常规场景下也是O(1)：loc为常数时间的结论适用于多级索引场景，只要设置的多级索引是唯一且已排序的（set_index默认会自动对索引排序），pandas会为索引构建哈希映射，查找时间复杂度为常数时间。但pandas的loc封装了大量额外逻辑，包括标签类型校验、切片兼容、数据对齐等，常数项开销远高于原生dict，实际查询速度比嵌套dict慢很多，小数据量下差距可达几十上百倍。
• 特殊场景下loc性能会退化：如果多级索引存在重复值，或者没有正常排序，loc查询可能会退化为*O(log n)的二分查找，甚至O(n)*的线性扫描，性能会大幅下降，而原生dict只要是正常使用就不会出现这类问题。

内容的提问来源于stack exchange，提问作者d.b