Python/Pandas 通过一个限制获得所有可能的组合

Question

目标 - 我总共有 50 条记录，需要找出工资 = 48,000 的 6 名球员的所有可能组合。

如果我只使用大约 20 条记录，下面的代码将起作用，但在尝试将其应用于所有 50 条记录时，我一直遇到内存错误。我正在寻找一种方法来优化我的代码，要么只接受 50k 以下的组合，而不会像我一样循环。

示例数据（目前共有 50 条记录）-

    ID          Salary
0   11282489    11000
1   11282517    10800
2   11282479    10700
3   11282521    10200
4   11282483    10100
5   11282481    10000

当前代码 -

comb = combinations(data['ID'], 6) 
comb_list = list(comb)
df_list = []
for i in comb_list:
    i = list(i)
    if data.loc[data['ID'].isin(i)]['Salary'].sum() <= 50000 and data.loc[data['ID'].isin(i)]['Salary'].sum() >= 48000:
        df_list.append(data.loc[data['ID'].isin(i)])

    counter +=1

“comb_list”目前以大约 1500 万个组合结束，这是主要问题。有没有比我现在做的更好的方法来应用薪水过滤器？

谢谢！

Answer 1

您当然可以避免循环。

找到所有组合，将它们的 ID 映射到薪水，然后计算每个组合的总和。然后只是子集到那些薪水在 48,000 到 50,000 之间的组合

设置

import pandas as pd
import numpy as np
from itertools import combinations

np.random.seed(123)
df = pd.DataFrame({'ID': np.arange(1,51,1),
                   'Salary': np.random.randint(7000,12000,50)})
# ID to Salary dictionary
d = df.set_index('ID').Salary.to_dict()

代码

n = 6  # length of combination tuples

# Create df of people and their salary
df2 = pd.DataFrame(list(combinations(df.ID, n)), 
                   columns=['p'+str(i) for i in np.arange(1,n+1,1)])
df2 = pd.concat([df2, df2.replace(d).add_suffix('_salary')], axis=1)

# Subset to those within the range you care about
df2[df2[[col for col in df2.columns if '_salary' in col]].sum(1).between(48000,50000)]

输出

        p1  p2  p3  p4  p5  p6  p1_salary  p2_salary  p3_salary  p4_salary  p5_salary  p6_salary
48465    1   2   6  10  19  32      10582      10454       7096       7111       7039       7588
48481    1   2   6  10  19  48      10582      10454       7096       7111       7039       7371
209845   1   3   5   6   9  10      10582       8346       8593       7096       7942       7111
209854   1   3   5   6   9  19      10582       8346       8593       7096       7942       7039
209883   1   3   5   6   9  48      10582       8346       8593       7096       7942       7371
...

（有 188,531 个这样的组合）。肯定会有更有效的解决方案。