熊猫在预算下找到行的所有组合

时间:2019-09-26 14:57:29

标签: python pandas dataframe itertools

我正在尝试找出一种方法来确定DataFrame中所有低于预算的行的组合,因此,假设我有一个像这样的数据框:

data = [['Bread', 9, 'Food'], ['Shoes', 20, 'Clothes'], ['Shirt', 15, 'Clothes'], ['Milk', 5, 'Drink'], ['Cereal', 8, 'Food'], ['Chips', 10, 'Food'], ['Beer', 15, 'Drink'], ['Popcorn', 3, 'Food'], ['Ice Cream', 6, 'Food'], ['Soda', 4, 'Drink']]
df = pd.DataFrame(data, columns = ['Item', 'Price', 'Type'])
df

数据

Item       Price  Type
Bread      9      Food
Shoes      20     Clothes
Shirt      15     Clothes
Milk       5      Drink
Cereal     8      Food
Chips      10     Food
Beer       15     Drink
Popcorn    3      Food
Ice Cream  6      Food
Soda       4      Drink

我想找到可以在特定预算下购买的每种组合,对于本示例来说,假设是35美元,而每种类型只得到其中一种。我想获得一个新的数据框,该数据框由适用于其自己列中每个项目的每种组合的行组成。

我试图使用itertools.product做到这一点,但这可以合并并添加列,但是我真正需要做的是根据另一列中的值合并并添加特定列。我现在有点难过。

感谢您的帮助!

1 个答案:

答案 0 :(得分:3)

这是一种将powerset中的itertools配方与pd.concat一起使用的方法

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))

df_groups = pd.concat([df.reindex(l).assign(grp=n) for n, l in 
                       enumerate(powerset(df.index)) 
                       if (df.loc[l, 'Price'].sum() <= 35)])

输出具有满足$ 35条件的产品组的单个数据框:

          Item  Price     Type  grp
0       Bread      9     Food    1
1       Shoes     20  Clothes    2
2       Shirt     15  Clothes    3
3        Milk      5    Drink    4
4      Cereal      8     Food    5
..        ...    ...      ...  ...
3        Milk      5    Drink  752
4      Cereal      8     Food  752
7     Popcorn      3     Food  752
8   Ice Cream      6     Food  752
9        Soda      4    Drink  752

通过几种方式可以满足35美元的预算?

df_groups['grp'].nunique()

输出:

258

详细信息:

这里使用了一些技巧/方法。首先,我们使用数据框的索引来使用powerset创建行或项的组。接下来,我们使用enumerate来标识每个组,并使用assign在数据框中使用枚举的组号创建一个新列。

修改以捕获每种类型的不超过一种:

df_groups = pd.concat([df.reindex(l).assign(grp=n) for n, l in 
                       enumerate(powerset(df.index)) 
                       if ((df.loc[l, 'Price'].sum() <= 35) & 
                           (df.loc[l, 'Type'].value_counts()==1).all())])

多少个小组?

df_groups['grp'].nunique()
62

为每种类型精确获取一个:

df_groups = pd.concat([df.reindex(l).assign(grp=n) for n, l in 
                       enumerate(powerset(df.index)) 
                       if ((df.loc[l, 'Price'].sum() <= 35) & 
                           (df.loc[l, 'Type'].value_counts()==1).all()&
                           (len(df.loc[l, 'Type']) == 3))])

多少个小组?

df_groups['grp'].nunique()
21