I have a data frame with column with numbers.

Now I want to pick the first 3 numbers if they match a certain condition

e.g. if column values start with 540 or 514 then we give me the first 3 numbers otherwise give the full number

def first_three(val):
    return val[:3] 
   if len(val) > 3 else val

But I want only the first 3 numbers only if they match the my specified condition

Thanks in advance

CodePudding user response：

def first_three(val):
    aaa = str(val)
    bbb = int(aaa[:3])
    qqq = 0
    if bbb == 514 or bbb == 540:
        qqq = bbb
    else:
        qqq = val
    return qqq 


print(first_three(514008))

The number is turned into a string, the first three elements are taken and converted back to int. Where in the verification process we choose which number to return.

CodePudding user response：

import pandas as pd
d = {'Year': [2022, 2022,2022],
     'Customer': [530010, 540060,514008],
    'Price':[11728.7,4793.21,-15665.40]}

df = pd.DataFrame(data=d)
df['new_columns']=[int(str(val)[0:3]) if (str(val)[0:3]=='540') or (str(val)[0:3]=='514') else val for val in df['Customer'].values  ]
df