How to concisely create new columns as output from a zip function?

user_10308255 0
Viewed 112

I have a dataframe I am adding new columns to. I am doing this using output from zip that uses output from a function. Seen below:

The function generates 4 new columns that I am trying to add to an existing dataframe.

data = [
    [1, 123],
    [2, 123454],
    [3, 64564],
]

df = pd.DataFrame(data, columns=["ID", "number"])

# function
def func(num):
    double = num * 2
    triple = num * 3
    quadruple = num * 4
    tenex = num * 10

    return double, triple, quadruple, tenex

# apply function to create 4 new columns
df["double"], df["triple"], df["quad"], df["tenex"] = zip(
    *df["number"].apply(lambda x: func(x))
)

Is there a more concise way to do this? It's fine when I am adding only 4 columns, but I want to expand this function to add 10+ columns.

I was considering something like this:

tuple(df[colname] for colname in col_list) = zip(
    *df["number"].apply(lambda x: func(x))
)

but it doesn't work (error: SyntaxError: cannot assign function to call)

pandas dataframe tuples zip python
4 Answers

I'd use a dictionary to map column names to functions - you don't really benefit from having all the computations done inside a single function.

something like:

column_mapper = {
    'double': lambda x: x*2,
    'triple': lambda x: x*3,
    'quadruple': lambda x: x*4,
}

data = [
    [1, 123],
    [2, 123454],
    [3, 64564],
]
df = pd.DataFrame(data, columns=["ID", "number"])

for column_name, func in column_mapper.items():
    df[column_name] = df['number'].apply(func)
user_111614320

Pass result_type='expand' to apply to output to multiple columns:

df[["double","triple","quad","tenex"]] = df.apply(lambda x: func(x['number']), axis = 1, result_type='expand')
user_114770310

I think the better way would actually be to create separate functions here. Create one function that takes two arguments, x and n, and then use functools.partial to create single-argument functions to use with apply:

from functools import partial


def multiply(x, n):
    return x * n


functions = ((col, partial(multiply, n=i)) for col, i in [('double', 2), ('triple', 3), ('quadruple', 4), ('tenx', 10)])

for col, func in functions:
    df[col] = df['number'].apply(func)
user_78679680

enumerate could work as long as the cols list in in the same order as tuple returned in func.

cols = ['double', 'triple', 'quad', 'tenex']

for i, col in enumerate(cols):
    df[col] = df["number"].apply(lambda x: func(x)[i])
user_151267630
Related