百万级重复随机整数列表的索引查找提速方案咨询

哥们，你这代码跑慢的核心问题太明显了——a.index(i)这个操作是线性遍历整个列表找元素，你有10000个左右的唯一元素，每个都要扫一遍100万条数据的列表，这相当于10000×100万=100亿次操作，不慢才怪！而且你代码里的inds=[0]和if i not in inds完全是多余的，unique_a是集合，元素不会重复，那行代码根本起不到任何作用。

下面给你几个针对性的优化方案，保证把运行时间从几分钟压到几秒甚至更短：

方案1：先记录原列表的首次索引，再排序唯一元素（对应你原代码的逻辑） #

这个方案是先遍历一次原列表，用字典把每个元素第一次出现的索引存下来（哈希表查找是O(1)的），之后直接从字典里取索引就行，不用再反复扫列表。

import random

# 生成列表：用列表推导式比for+append更快，底层优化过
a = [random.randint(1, 10000) for _ in range(1000000)]

# 只遍历一次原列表，记录每个元素的首次出现索引
first_occurrence = {}
for idx, num in enumerate(a):
    # 只在元素没被记录过时才赋值，保证是首次出现的索引
    if num not in first_occurrence:
        first_occurrence[num] = idx

# 对唯一元素排序，然后提取对应的首次索引
sorted_unique_nums = sorted(first_occurrence.keys())
inds = [first_occurrence[num] for num in sorted_unique_nums]

这个方案的时间复杂度是O(n + m log m)，n是100万，m是1万左右，跑起来绝对秒出结果。

方案2：如果是要排序后的列表中首次出现的索引（对应你需求描述的“先排序再找索引”） #

如果你的真实需求是：先把整个列表排序，然后在排序后的列表里找每个唯一元素第一次出现的位置，那可以这么做：

方法A：遍历排序后的列表记录索引 #

import random

a = [random.randint(1, 10000) for _ in range(1000000)]
sorted_a = sorted(a)

first_indices = {}
for idx, num in enumerate(sorted_a):
    if num not in first_indices:
        first_indices[num] = idx

# 按元素排序后提取索引
sorted_unique_nums = sorted(first_indices.keys())
inds = [first_indices[num] for num in sorted_unique_nums]

方法B：用itertools.groupby跳过重复元素 #

排序后相同元素是连续的，用groupby可以直接跳过整个重复组，减少遍历次数：

from itertools import groupby
import random

a = [random.randint(1, 10000) for _ in range(1000000)]
sorted_a = sorted(a)

first_indices = {}
current_idx = 0
for num, group in groupby(sorted_a):
    first_indices[num] = current_idx
    # 直接跳到下一个不同元素的位置
    current_idx += len(list(group))

sorted_unique_nums = sorted(first_indices.keys())
inds = [first_indices[num] for num in sorted_unique_nums]

方案3：用numpy直接开挂（最快的方式） #

如果可以用numpy库，那直接用它的内置函数，底层是C实现的，速度碾压纯Python代码：

import numpy as np

# 生成百万级随机整数数组
a = np.random.randint(1, 10001, size=1000000)
# np.unique的return_index参数直接返回每个唯一元素在原数组的首次索引
unique_nums, first_indices = np.unique(a, return_index=True)
# 对唯一元素排序，然后对应提取索引
sorted_order = np.argsort(unique_nums)
sorted_unique_nums = unique_nums[sorted_order]
inds = first_indices[sorted_order].tolist()

这个方案跑起来基本是瞬间完成，比纯Python快好几倍。

最后再提几个小细节： #

1. 生成列表时，用列表推导式[... for _ in range(...)]比for x in range(...): a.append(...)更快，因为列表推导式是Python底层优化过的。
2. 永远避免在循环里用list.index()这种线性查找，哈希表（字典）是解决这类问题的最优选择。

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