如何以Pythonic方式获取指定日期后的下一个1月、5月、9月？

Your existing solution works, but generating a DatetimeIndex and scanning through it isn't the most efficient or readable way to solve this problem. Let's replace that with a direct computation that avoids unnecessary date range generation and leans into Python's concise syntax.

Here's a refined version:

import pandas as pd

def rel_month(dt):
    ref_dt = pd.Timestamp(dt)
    target_months = [1, 5, 9]
    
    # Filter target months that come after the reference month in the same year
    upcoming_in_year = [month for month in target_months if month > ref_dt.month]
    
    if upcoming_in_year:
        # Pick the earliest upcoming month in the same year
        next_month = min(upcoming_in_year)
        next_year = ref_dt.year
    else:
        # No target months left this year; use the first one of the next year
        next_month = target_months[0]
        next_year = ref_dt.year + 1
    
    # Format and return the result
    return pd.Timestamp(f"{next_year}-{next_month:02d}").strftime('%Y-%m')

Why this is better: #

• Efficiency: No need to generate a range of dates (the 122D window was a safe guess but unnecessary). We directly compute the result in constant time.
• Readability: The logic is explicit—anyone reading the code can immediately follow how we're finding the next target month.
• Pythonic: Uses list comprehensions (a core Python idiom) instead of clunky index scanning, and avoids explicit for loops with manual iteration.

Test cases to verify: #

• rel_month("2016-02") returns 2016-05 (matches your example)
• rel_month("2016-10") returns 2017-01
• rel_month("2016-05") returns 2016-09
• rel_month("2016-09") returns 2017-01

If you need to handle a pandas Series of dates instead of a single date, here's a vectorized version that scales efficiently:

def rel_month_vectorized(dates):
    ref_dts = pd.to_datetime(dates)
    target_months = pd.Series([1,5,9])
    
    def get_next_month(month):
        mask = target_months > month
        return target_months[mask].min() if mask.any() else target_months[0]
    
    next_months = ref_dts.dt.month.apply(get_next_month)
    next_years = ref_dts.dt.year + (next_months < ref_dts.dt.month).astype(int)
    
    return pd.to_datetime(dict(year=next_years, month=next_months, day=1)).dt.strftime('%Y-%m')

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