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

Question

目标-我总共有50条记录，需要找到6位薪水<= 50,000和> = 48,000的球员的所有可能组合。

如果我仅使用大约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个这样的组合）。必然会有更有效的解决方案。