如何基于滚动窗口条件创建二元变量？

原代码的核心问题 #

1. 变量名错误：代码中引用了不存在的列e_polity2，数据源的分数列实际是score，这会直接触发KeyError。
2. 窗口逻辑不匹配：你需要的是非重叠的3年固定窗口（如1950-1952、1953-1955），但原代码使用了滑动窗口rolling(3)，且仅检查相邻年份的分数差，无法精准匹配你的分组需求。

方案1：非重叠3年固定窗口（匹配你给出的示例规则） #

该方案将数据按每3年一组划分，只要组内任意两年间分数提升≥2，组内所有年份的二元变量都设为1。

import pandas as pd

data = {
    'year': [1950, 1951, 1952, 1953, 1954, 1955, 1956, 1957, 1958],
    'teamcode': [4, 4, 4, 4, 4, 4, 4, 4, 4],
    'teamname': ['Team 1'] * 9,
    'score': [0, 1, 5, 5, 5, 6, -10, -10, -10]
}
df = pd.DataFrame(data)

# 生成非重叠3年窗口分组：1950-52为组0，1953-55为组1，1956-58为组2
df['window_group'] = (df['year'] - df['year'].min()) // 3

# 检查分组内是否存在任意两年分数提升≥2
def check_improvement(group):
    scores = group['score'].values
    for i in range(len(scores)):
        for j in range(i+1, len(scores)):
            if scores[j] - scores[i] >= 2:
                return 1
    return 0

# 计算每个分组的标记并合并回原数据
group_marks = df.groupby('window_group').apply(check_improvement).reset_index(name='group_flag')
df = df.merge(group_marks, on='window_group')
df['binary_var'] = df['group_flag']

# 清理临时列
df.drop(['window_group', 'group_flag'], axis=1, inplace=True)

print(df)

运行结果 #

year  teamcode teamname  score  binary_var
0  1950         4   Team 1      0           1
1  1951         4   Team 1      1           1
2  1952         4   Team 1      5           1
3  1953         4   Team 1      5           0
4  1954         4   Team 1      5           0
5  1955         4   Team 1      6           0
6  1956         4   Team 1    -10           0
7  1957         4   Team 1    -10           0
8  1958         4   Team 1    -10           0

方案2：滑动3年窗口（若需连续滑动的窗口规则） #

如果需要连续滑动的3年窗口（如1950-52、1951-53、1952-54...），只要年份所在的任意一个窗口满足条件，二元变量设为1：

import pandas as pd

data = {
    'year': [1950, 1951, 1952, 1953, 1954, 1955, 1956, 1957, 1958],
    'teamcode': [4, 4, 4, 4, 4, 4, 4, 4, 4],
    'teamname': ['Team 1'] * 9,
    'score': [0, 1, 5, 5, 5, 6, -10, -10, -10]
}
df = pd.DataFrame(data)

# 检查滑动窗口内是否存在符合条件的分数提升
def check_sliding_window(window):
    scores = window.values
    for i in range(len(scores)):
        for j in range(i+1, len(scores)):
            if scores[j] - scores[i] >= 2:
                return 1
    return 0

# 生成每个滑动窗口的标记
df['window_flag'] = df['score'].rolling(window=3, min_periods=3).apply(check_sliding_window, raw=True)

# 将窗口标记向前填充，覆盖窗口内的所有年份
df['binary_var'] = df['window_flag'].rolling(window=3, min_periods=1).max().fillna(0).astype(int)

# 清理临时列
df.drop('window_flag', axis=1, inplace=True)

print(df)

运行结果 #

year  teamcode teamname  score  binary_var
0  1950         4   Team 1      0           0
1  1951         4   Team 1      1           0
2  1952         4   Team 1      5           1
3  1953         4   Team 1      5           1
4  1954         4   Team 1      5           1
5  1955         4   Team 1      6           0
6  1956         4   Team 1    -10           0
7  1957         4   Team 1    -10           0
8  1958         4   Team 1    -10           0

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