MemoryError pandas根据其他列值生成新列

Question

I have a data frame. structure as following:

OMT                              object
ZIPCODE                          object
PRODUCT_CAT                       int64
SERVICE_CATEGORY                 object
CURRENT_STANDARD_EDD            float64
TOTAL                             int64
DESTINATION_DISTRIBUTION_CTR     object
OPS_EDD                         float64
OPS_EDD_achieve                   int64
suggest_edd_1                    object
suggest_edd_2                     int64
suggest_edd_value_1               int64
suggest_edd_value_2               int64
final_edd_group                  object
final_edd                       float64
final_edd_value                   int64

I want to perform the following operation: when the total is < 5, return the label value of among D1/D2/D3/D4/D5/D6 where the first value over -1 compared with D6.(if none, D6)

if total is >=5, return the label value of among D1/D2/D3/D4/D5/D6 where the first value / d5 over 0.95 compared with D6.(if none, D6)

I wrote following code but return

 training_group['suggest_edd_1'] =np.where(training_group['TOTAL']>5,training_group[['D1','D2',
                'D3','D4','D5',
                'D6']].sub(training_group['D6'],axis =0).ge(-1).assign(D6=True).idxmax(1).str.extract('(\d+)'),
                 training_group[['D1','D2',
                'D3','D4','D5',
                'D6']].div(training_group['TOTAL'],axis =0).ge(0.95).assign(D6=True).idxmax(1).str.extract('(\d+)')) 

<ipython-input-72-61626eae2be9> in <module>
      4                  training_group[['D1','D2',
      5                 'D3','D4','D5',
----> 6                 'D6']].div(training_group['TOTAL'],axis =0).ge(OD_pari_target).assign(D6=True).idxmax(1).str.extract('(\d+)')) 

MemoryError: 

(Each separate one works, but if I apply condition on TOTAL it does not work.

I tried to use the lambda function applies to each row, but I could not find the approriate code to replace

assign(D6=True) and extract function

    if x['TOTAL'] < piece_threthold:
        return x[['D1','D2',
                'D3','D4','D5',
                'D6']].sub(x['D6'],axis =0).ge(OD_pari_piece).ge(-1).idxmax(1)
    else:
        return x[['D1','D2',
                'D3','D4','D5',
                'D6']].div(x['TOTAL'],axis =0).ge(OD_pari_target).ge(-1).idxmax(1)

I can get the result I need by doing the following. however, I feel it is very inefficient and create more columns which I do not need. (I will drop the suggest_edd_1 and suggest_edd_2 later since I only need the final_suggest)

training_group['suggest_edd_1'] = training_group[['D1','D2',
                'D3','D4','D5',
                'D6']].sub(training_group['D6'],axis =0).ge(OD_pari_piece).assign(D6=True).idxmax(1).str.extract('(\d+)')

training_group['suggest_edd_2'] = training_group[['D1','D2',
                'D3','D4','D5',
                'D6']].div(training_group['TOTAL'],axis =0).ge(OD_pari_target).assign(D6=True).idxmax(1).str.extract('(\d+)')

training_group['final_suggest'] = np.where(training_group['TOTAL']>5,training_group['suggest_edd_1'] ,training_group['suggest_edd_2'])

Answer 1

As you side each of them work fine pre-calculated the value to assign

s1=training_group[['D1','D2',
                'D3','D4','D5',
                'D6']].sub(training_group['D6'],axis =0).ge(-1).assign(D6=True).idxmax(1).str.extract('(\d+)')
s2=training_group[['D1','D2',
                'D3','D4','D5',
                'D6']].div(training_group['TOTAL'],axis =0).ge(0.95).assign(D6=True).idxmax(1).str.extract('(\d+)')

training_group['suggest_edd_1'] =np.where(training_group['TOTAL']>5,s1,
                 s2)