使用pandas按指定规则将数据框中水果分配到对应日期的实现方案

需求规则 #

核心分配规则 #

• 规则1：apple不可连续发放，两个apple的发放间隔至少为1天
• 规则2：orange必须在第7天（最后一天）发放
• 规则3：第6天不可发放lemon，且lemon与orange的发放间隔至少为1天
• 规则4：lemon和orange的发放日不可搭配其他水果；若总水果数超过7，多余水果可与apple、cherry、banana同天发放

测试场景数据 #

基础场景（总水果数=7）

输入数据：

fruit  count   
apple  3  
orange 1  
lemon  1  
cherry 1  
banana 1

预期输出：

day fruit
d1   apple
d2   cherry
d3   apple
d4   lemon
d5   apple
d6   banana
d7   orange

场景1（总水果数<7）

输入数据：

fruit  count   
apple  1  
orange 1  
lemon  1  
cherry 1  
banana 1

预期输出：

day fruit
d1   apple
d2   cherry
d4   lemon
d6   banana
d7   orange

场景2（总水果数>7）

输入数据：

fruit  count  
apple  3  
orange 1
pomegranate 1  
lemon  1  
cherry 1  
banana 1 

预期输出：

day fruit
d1   apple
d2   cherry
d3   apple
d3   pomegranate
d4   lemon
d5   apple
d6   banana
d7   orange

调整后实现代码 #

import pandas as pd

def allocate_fruits(fruit_df):
    # 初始化7天存储结构
    day_map = {f'd{i}': [] for i in range(1, 8)}
    # 规则2：固定第7天放orange
    day_map['d7'].append('orange')
    fruit_cnt = fruit_df.set_index('fruit')['count'].to_dict()
    
    # 第一步：分配apple，满足间隔要求，避开第7天
    apple_num = fruit_cnt.get('apple', 0)
    available_apple_days = ['d1', 'd3', 'd5']
    for i in range(apple_num):
        if i < len(available_apple_days):
            day_map[available_apple_days[i]].append('apple')
    
    # 第二步：分配lemon，满足规则3、4
    lemon_num = fruit_cnt.get('lemon', 0)
    if lemon_num > 0:
        # 筛选符合要求的空日期：不能是d6、d7，当天无其他水果
        available_lemon_days = [d for d in day_map.keys() if d not in ['d6', 'd7'] and len(day_map[d]) == 0]
        if available_lemon_days:
            # 选最靠后的日期保证和orange间隔足够
            day_map[available_lemon_days[-1]].append('lemon')
    
    # 第三步：分配其他普通水果
    normal_fruits = []
    for fruit, cnt in fruit_cnt.items():
        if fruit not in ['apple', 'orange', 'lemon']:
            normal_fruits.extend([fruit]*cnt)
    
    # 优先填充空白日期
    empty_days = [d for d in day_map.keys() if len(day_map[d]) == 0]
    for idx, fruit in enumerate(normal_fruits):
        if idx < len(empty_days):
            day_map[empty_days[idx]].append(fruit)
        else:
            # 多余水果放到允许混放的日期
            for d in day_map.keys():
                if len(set(day_map[d]) & {'lemon', 'orange'}) == 0:
                    day_map[d].append(fruit)
                    break
    
    # 转换为pandas DataFrame输出
    res = []
    for day, fruits in day_map.items():
        for f in fruits:
            res.append({'day': day, 'fruit': f})
    return pd.DataFrame(res)

# 测试示例，替换fruit_df即可适配不同场景
fruit_df = pd.DataFrame({
    'fruit': ['apple', 'orange', 'lemon', 'cherry', 'banana', 'pomegranate'],
    'count': [3, 1, 1, 1, 1, 1]
})
print(allocate_fruits(fruit_df))

内容的提问来源于stack exchange，提问作者rakesh