Как добавить дополнительный столбец в dataframe с проверкой условия?

Question

У меня проблема с разбивкой информационного кадра на две строки. У меня есть следующие два кадра данных:

df1 = pd.DataFrame({'Item':['A','B','C'],'Value':[4,3,7]})
df2 = pd.DataFrame({'Item':['A','B','C'],'Value':[4,3,2]})

df1:
    Item    Value
0      A        4
1      B        3
2      C        7

df2:
    Item    Value
0      A        4
1      B        3
2      C        2

Я хотел бы установить предел для значения, скажем, допустим, предел = 10

limit = 10

После этого я хочу добавить дополнительный столбец и проверим, находится ли товар в пределах лимита, если нет, я хочу разбить его на две строки и показать следующий результат: Однако, если он находится в пределах лимита, я хочу добавить дополнительную строку, чтобы показать, сколько осталось денег на предел:

df1:
    Item    Value   within_limit
0      A        4            Yes
1      B        3            Yes
2      C        3            Yes
3      C        4             No

df2:
    Item    Value   within_limit
0      A        4            Yes
1      B        3            Yes
2      C        2            Yes
3    Nan        1          Extra

Как я могу получить вышеуказанный результат?

Answer 1

Вот ваша функция, предполагается, что индекс является индексом диапазона:

def check_limit(df, limit):
    s = df['Value'].cumsum()
    df['within_limit'] = np.where(s.le(limit), 'Yes', 'No')

    # Nothing off limit
    if s.iloc[-1] < limit:
        return df.append(pd.Series([np.nan, limit-s.iloc[-1],'Extra'], 
                                  index=df.columns,
                                  name=len(df))
                      )

    # all sum even to limit
    if s.iloc[-1] == limit: return df

    # find where the limit is exceeded
    idx = s.gt(limit).idxmax()
    # exceed limit
    exceed_limit = s.loc[idx] - limit

    new_df = df.loc[[idx,idx]].copy()
    new_df['Value'] = [df.loc[idx,'Value'] - exceed_limit, exceed_limit]
    new_df['within_limit'] = ['Yes','No']

    return pd.concat((df.drop(idx), new_df)).sort_index().reset_index(drop=True)

# test data
df1 = pd.DataFrame({'Item':['A','B','C', 'D'],'Value':[4,3,7,2]})
df2 = pd.DataFrame({'Item':['A','B','C'],'Value':[4,3,2]})

Вывод:

# check_limit(df1, 10)
  Item  Value within_limit
0    A      4          Yes
1    B      3          Yes
2    C      3          Yes
3    C      4           No
4    D      2           No

# check_limit(df2, 10)
  Item  Value within_limit
0    A      4          Yes
1    B      3          Yes
2    C      2          Yes
3  NaN      1        Extra