Python从含缺失元素的季度/月份数组构建映射字典的优化问询

基础问题 #

给定季度列表和月份列表：

quarters = ['Q1-21', 'Q2-21', 'Q3-21', 'Q4-21', 'Q1-22', 'Q2-22', 'Q3-22', 'Q4-22']  

months = [
    'JAN-21', 'FEB-21', 'MAR-21', 'APR-21', 'MAY-21', 'JUN-21', 
    'JUL-21','AUG-21', 'SEP-21','OCT-21', 'NOV-21', 'DEC-21', 
    'JAN-22', 'FEB-22','MAR-22', 'APR-22', 'MAY-22', 'JUN-22', 
    'JUL-22', 'AUG-22', 'SEP-22','OCT-22', 'NOV-22', 'DEC-22'
]

需要创建一个以季度为键、对应月份列表为值的字典。当两个列表完整且有序时，可通过拆分月份数组并与季度列表配对实现：

import numpy as np
dict(zip(quarters, np.split(months, 8)))

输出结果：

{'Q1-21': ['JAN-21', 'FEB-21', 'MAR-21'],
 'Q2-21': ['APR-21', 'MAY-21', 'JUN-21'],
 'Q3-21': ['JUL-21', 'AUG-21', 'SEP-21'],
 'Q4-21': ['OCT-21', 'NOV-21', 'DEC-21'],
 'Q1-22': ['JAN-22', 'FEB-22', 'MAR-22'],
 'Q2-22': ['APR-22', 'MAY-22', 'JUN-22'],
 'Q3-22': ['JUL-22', 'AUG-22', 'SEP-22'],
 'Q4-22': ['OCT-22', 'NOV-22', 'DEC-22']}

扩展问题 #

当季度或月份列表存在元素缺失时（如下示例），需要适配该场景：

quarters = ['Q2-21', 'Q3-21', 'Q4-21', 'Q1-22', 'Q2-22', 'Q3-22']  
months = [
    'MAY-21', 'JUN-21', 
    'JUL-21','AUG-21', 'SEP-21',
    'OCT-21', 'NOV-21', 
    'JAN-22', 'FEB-22','MAR-22', 
    'APR-22', 'MAY-22', 'JUN-22', 
    'JUL-22', 'AUG-22', 'SEP-22',
    'OCT-22'
]

期望输出为：

{'Q2-21': ['MAY-21', 'JUN-21'],
 'Q3-21': ['JUL-21', 'AUG-21', 'SEP-21'],
 'Q4-21': ['OCT-21', 'NOV-21'],
 'Q1-22': ['JAN-22', 'FEB-22', 'MAR-22'],
 'Q2-22': ['APR-22', 'MAY-22', 'JUN-22'],
 'Q3-22': ['JUL-22', 'AUG-22', 'SEP-22'],
 'Q4-22': ['OCT-22']}

初始解决方案 #

通过datetime和偏移量生成每个季度对应的月份字符串，再筛选出存在于月份列表中的项，最后构建字典：

import calendar
import pandas as pd

month_vals = [
    [
        # 生成对应的月份字符串
        calendar.month_name[m.month][:3].upper() + q[-3:]  
        # 列出季度内的所有月份
        for m in pd.date_range(pd.to_datetime(q[-2:] + q[:2]), 
                               pd.to_datetime(q[-2:] + q[:2]) + pd.offsets.QuarterEnd(),
                               freq='MS')
        # 检查是否在月份列表中
        if calendar.month_name[m.month][:3].upper() + q[-3:] in months
    ] for q in quarters
]

# 创建字典
{k: v for k, v in zip(quarters, month_vals)}

该方案需要多次生成日期并做成员判断，效率偏低，需更优雅的实现方式。

优化实现 #

可以先建立月份到季度的映射规则，通过一次遍历完成月份分组，再整理出符合要求的字典，无需依赖第三方库，效率更高：

代码实现 #

# 定义月份缩写到季度的映射
month_to_quarter = {
    'JAN': 'Q1', 'FEB': 'Q1', 'MAR': 'Q1',
    'APR': 'Q2', 'MAY': 'Q2', 'JUN': 'Q2',
    'JUL': 'Q3', 'AUG': 'Q3', 'SEP': 'Q3',
    'OCT': 'Q4', 'NOV': 'Q4', 'DEC': 'Q4'
}

# 遍历月份列表，完成分组
quarter_groups = {}
for month in months:
    mon_abbr, year = month.split('-')
    q_key = f"{month_to_quarter[mon_abbr]}-{year}"
    if q_key not in quarter_groups:
        quarter_groups[q_key] = []
    quarter_groups[q_key].append(month)

# 整合所有需要的季度：包括输入的quarters和存在月份的季度
all_target_quarters = set(quarters) | set(quarter_groups.keys())
# 按年份和季度排序，保证输出顺序合理
sorted_quarters = sorted(all_target_quarters, key=lambda x: (x[-2:], int(x[1])))

# 构建最终字典
result = {q: quarter_groups.get(q, []) for q in sorted_quarters}

print(result)

输出结果 #

{'Q2-21': ['MAY-21', 'JUN-21'],
 'Q3-21': ['JUL-21', 'AUG-21', 'SEP-21'],
 'Q4-21': ['OCT-21', 'NOV-21'],
 'Q1-22': ['JAN-22', 'FEB-22', 'MAR-22'],
 'Q2-22': ['APR-22', 'MAY-22', 'JUN-22'],
 'Q3-22': ['JUL-22', 'AUG-22', 'SEP-22'],
 'Q4-22': ['OCT-22']}

方案优势 #

• 轻量高效：无需依赖pandas或datetime处理，仅遍历月份列表一次，时间复杂度为O(n)（n为月份数量）；
• 逻辑清晰：通过静态映射规则完成分组，代码易读易维护；
• 兼容性强：自动适配季度或月份缺失的场景，同时保证输出顺序符合时间逻辑。

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