如何加快在数据框列中查找重复项

Question

我希望找到数据框列中出现重复值序列的索引。我希望结果是一个列表列表，其中每个子列表都是一个单独的重复值索引序列。

我当前的代码可以工作，但速度很慢（在 10,000 行数据帧中 10% 的重复需要 15 毫秒）：

import pandas as pd
import numpy as np
import time

# Given a dataframe and column, return a list of lists where each sublist
# contains indexes of the sequential duplicates
def duplicate_ranges(df, c):
    return to_ranges(df[c].shift(1) == df[c])

# Take a pandas Series of booleans and return a list of lists where each 
# sub-list is the indexs of sequential true values in the list
def to_ranges(s):
    r = []
    g = []
    for k, v in s.items():
        if v == True:
            g.append(k)
        elif len(g) > 0:
            r.append(g)
            g = []
    if len(g) > 0:
        r.append(g)
    return r

def bench_it(n):
    data = {"A": np.random.randint(10, 10000)}
    idxs = pd.date_range(start='2000-01-01', periods=10000)
    df = pd.DataFrame(data, index=idxs)
    t = time.time()
    for _ in range(0, n):
        r = duplicate_ranges(df, 'A')
    t = time.time() - t
    print("{:d} iterations took {:.1f} msec".format(n, 1000*t))

bench_it(1000)

据我所知，所有时间都花在了 to_ranges() 的主循环中。我对 pandas 和 numpy 还很陌生，有人可以提出一种加快速度的方法吗？

Answer 1

这是一种利用scipy.sparse 内高效操作的快速方法：

from scipy.sparse import csr_matrix

def duplicate_ranges(df, c):
    index, values = df.index.values, df[c].values

    data = values
    indices = np.arange(len(values))
    indptr = np.concatenate([[0], np.where(np.diff(values) != 0)[0] + 1,
                             [len(values)]])

    M = csr_matrix((index, indices, indptr))[np.diff(indptr) > 1]
    M.sort_indices()
    return np.split(M.data, M.indptr[1:-1])

它比这里的其他方法快大约一个数量级，因为它避免了整个数组上的 Python 循环（尽管split() 函数中有一些 Python 循环，只在数组的子集上调用）。

---

旧答案：

这是解决此问题的快速方法：

df = pd.DataFrame({'A': [1, 2, 3, 3, 3, 2, 1, 1, 2, 2]})

def duplicate_ranges(df, c):
    index, values = df.index.values, df[c].values
    ranges = np.split(index, np.where(np.diff(values) != 0)[0] + 1)
    return [list(r) for r in ranges if len(r) > 1]

duplicate_ranges(df, 'A')
# [[2, 3, 4], [6, 7], [8, 9]]

由于它避免了嵌套循环并且只需要一次通过整个列，因此它应该比其他方法快得多。