Question

isSpring = False

# Calculate season e.g. spring wheat season is only ~ May - August
if isSpring:
    # Change from name of month to a number e.g. January is 1, December is 12
    idx_srt_mon = strptime('May', '%B').tm_mon
    idx_end_mon = strptime('September', '%B').tm_mon
else:
    idx_srt_mon = strptime('September', '%B').tm_mon
    idx_end_mon = strptime('August', '%B').tm_mon

# Assign seasons, works for crops with season within 1 year and also winter crops with seasons spanning years
df['Season'] = numpy.nan
first_yr     = df.YEAR.unique()[0]
for i,  row in df.iterrows():
    cur_month = row.MONTH
    cur_yr    = row.YEAR

    if cur_month < idx_srt_mon:
        if not isSpring and (cur_month < idx_end_mon):
            if cur_yr > first_yr:
                df.set_value(i, 'Season', cur_yr - 1)
    else:
        if cur_month <= idx_end_mon or not isSpring:
            df.set_value(i, 'Season', cur_yr)

在此处提供的数据框中：https://www.dropbox.com/s/gwyim6kc83mw3yf/df.csv?dl=0，我想添加一个新列，分配如下“季节”：

如果isSpring为TRUE，则将当前年份作为季节分配给从idx_srt_mon到idx_end_mon（5月到9月）的所有行
如果isSpring是假的话，那么这个季节从当年的9月开始，一直到明年8月。整个时期，跨越2年，这个季节是它开始的一年。例如，如果开始时间是2000年9月，则该季节将是2000年9月至2001年8月的2000年。
对于所有其他情况，季节是numpy.nan

我使用df.iterrows来做这个计算，但它很慢。我怎样才能加快速度呢？

Answer 1

您需要仔细检查逻辑，但可以通过构建正确的布尔条件来对其进行矢量化，如下所示。

In [20]: first_yr = df.YEAR.min()

In [21]: if isSpring:
    ...:     df.loc[df['MONTH'] <= 9, 'Season'] = df.YEAR
    ...: else:
    ...:     df['Season'] = np.where(
                 (df['MONTH'] < 8) & (df['YEAR'] > first_yr), 
                     df['YEAR'] - 1, 
                     np.where(df['MONTH'] <= 8, df['YEAR'], 
                                                np.nan))

加快大熊猫对季节的计算

1 个答案: