Python 中的 Input() IF ELSE 控制语句

Question

我有一个数据集，只要输入日期值，特定日期列就会发生变化。

数据

location    type    mig1        de          mig2        re
ny          aa      8/1/2021    10/1/2021   1/1/2022    2/1/2022
ny          aa      8/1/2021    10/1/2021   1/1/2022    2/1/2022
ca          aa      8/1/2021    10/1/2021   1/1/2022    2/1/2022
tx          bb      9/1/2021    11/1/2021   2/1/2022    3/1/2022

希望

 #Date 8/1/2022 is input, which yields:

    location    type    mig1        de          mig2        re
    ny          aa      03/1/2022   05/1/2022   8/1/2022    7/1/2022
    ny          aa      03/1/2022   05/1/2022   8/1/2022    7/1/2022
    ca          aa      03/1/2022   05/1/2022   8/1/2022    7/1/2022
    tx          bb      03/1/2021   11/1/2021   08/1/2022   3/1/2022

输入提示将询问用户他们希望输入的日期值。 用户输入日期“2021 年 8 月 1 日”，会根据上述规则更新剩余列的日期值。。

正在做

datevalue = pd.to_datetime(input("Enter shift: "))


 if(df[type] == 'aa'):

         d = {
             'mig1': pd.DateOffset(months=5),
             'de': pd.DateOffset(months=3),
             're': pd.DateOffset(months=1),
             }
         s = pd.Series(d).rsub(datevalue)
         df.assign(**{**s, 'mig2': datevalue})
 else:
         e = {
             'mig1': pd.DateOffset(months=5),
        
             }
         s1 = pd.Series(e).rsub(datevalue)
         df.assign(**{**s1, 'mig2': datevalue})

逻辑：

mig2 列中的日期值输入到 input() 提示符中，mig1、de 和 re 值根据以下规则移动：

   mig1 is 5 months from the date entered in the input() prompt
   de   is 3 months from the date entered in the input() prompt
   re   is 1 month from the date entered in the input() prompt
   Where [type] column values == 'aa'
   

唯一发生变化的数据是日期，它们基​​本上会根据用户输入发生变化。

欢迎提出任何建议。

Answer 1

IIUC：

input = datetime(2022, 8, 1)
conditions = {'mig1':5, 'de':3, 're':1}
def apply_this(x):
    if x.name == 'mig2':
        return [input]*len(x)
    else:
        return [input - pd.DateOffset(months=conditions[x.name])]*len(x)

date_cols =['mig1', 'de', 're', 'mig2']
df.loc[df['type'] == 'aa', date_cols] = df[date_cols].apply(lambda x: apply_this(x))

print(df)

OUTPUT

  location type       mig1         de       mig2         re
0       ny   aa 2022-03-01 2022-05-01 2022-08-01 2022-07-01
1       ny   aa 2022-03-01 2022-05-01 2022-08-01 2022-07-01
2       ca   aa 2022-03-01 2022-05-01 2022-08-01 2022-07-01
3       tx   bb 2021-09-01 2021-11-01 2022-02-01 2022-03-01

SETUP

data = {'location': ['ny', 'ny', 'ca', 'tx'], 
        'type': ['aa', 'aa', 'aa', 'bb'], 
        'mig1': ['2021-08-01 00:00:00', '2021-08-01 00:00:00', '2021-08-01 00:00:00', '2021-09-01 00:00:00'],
        'de': ['2021-10-01 00:00:00', '2021-10-01 00:00:00', '2021-10-01 00:00:00', '2021-11-01 00:00:00'],
        'mig2': ['2022-01-01 00:00:00', '2022-01-01 00:00:00', '2022-01-01 00:00:00', '2022-02-01 00:00:00'],
        're': ['2022-02-01 00:00:00', '2022-02-01 00:00:00', '2022-02-01 00:00:00', '2022-03-01 00:00:00']}


df = pd.DataFrame(data)
for col in ['mig1', 'de', 'mig2', 're']:
    df[col] = pd.to_datetime(df[col])

根据 cmets 中的讨论进行编辑：

conditions = {'de':3, 're':1}
def apply_this(x):
    return [input - pd.DateOffset(months=conditions[x.name])]*len(x)

date_cols =['de', 're']
df.loc[df['type'] == 'aa', date_cols] = df[date_cols].apply(lambda x: apply_this(x))
df['mig1'] = input - pd.DateOffset(months=5)
df['mig2'] = input