找出 pandas dataframe 列中数字的所有可能组合，总和为 0

Question

我有一个 dataframe 有 10000 行..我想找到特定列（数量）中的值总和为 0 的所有行组合。

df=

ID_Key          Amount
10               12.4
12               -26.6
13                14.2        
14                15
17                 4.5
18                -9
19                94
20                -6

结果 dataframe 是

Combinations        Sum
(10,12,13)            0
(14,18,20)            0

下面是所有 3 个数字的组合代码，总和为 0。我也必须写 4 个数字和 5 个数字的组合，总和为 0，但即使是 3 个数字，当 dataframe 大小超过 30 时，它变得非常慢。如何降低以下算法的时间复杂度

from itertools import combinations
lst = [] 
t_counter=0
#all combinations ID_key consisting of length 3
for tuple_nums in set(combinations(df['ID_Key'], 3)):

    if df.shape[0]>2:
        t_counter=t_counter+1
        if df.loc[df['ID_Key'].isin(tuple_nums)].empty==False:
            if df.loc[df['ID_Key'].isin(tuple_nums), 'Amount'].sum()==0:
                lst.append([tuple_nums,df.loc[df['ID_Key'].isin(tuple_nums), 'Amount'].sum()])
                df=df.loc[~df['ID_Key'].isin(tuple_nums)]


    else:
        break



df_final=pd.DataFrame(lst, columns=['Combinations', 'Sum'])

我认为问题之一是在下面的代码中，迭代器 tuple_nums 遍历for tuple_nums in set(combinations(df['ID_Key'], 3))所有可能组合我每次都减少 dataframe 的大小，我得到的组合是在这一行中总和为 0 df=df.loc[~df['ID_Key'].isin(tuple_nums)]但仍然会遍历所有可能的组合。如何降低时间复杂度并使算法更快地处理 10000行

Answer 1

您可以使用几行代码来完成此操作，但计算量会很大。

让我们使用itertools中的powerset配方：

from itertools import chain, combinations
def powerset(iterable):
    "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
    s = list(iterable)
    return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))

然后像这样使用列表理解：

[tuple(i) for i in powerset(df.index) if len(i) > 0 and df.loc[list(i), 'Amount'].sum() == 0]

Output：

[(14, 18, 20)]

请注意，不确定您是如何从给定的数据中得到 (10,12,13) 的。