0
Follow
0
View

How to speed up nested loop and add condition?

dishaoye_521 注册会员
2023-01-25 16:34

You could use multiprocessing if you have more than one CPU.

from multiprocessing import Pool

def get_customer_data(data_tuple) -> dict:
    results = {}
    customer, state, amount, start, group = data_tuple
    for year in range(start, len(multi_data)+1):
        if year == start:
            results[customer] = [[amount * multi_data[year-1, state, :]]] 
        else:
            results[customer].append([results[customer][-1][-1] @ multi_data[year-1]])

        for lookup in lookups:
            if isinstance(lookup.index, pd.MultiIndex):
                value = lookup.loc[(year, group)].iat[0]
            else:
                value = lookup.loc[year].iat[0]
            results[customer][-1].append(value * results[customer][-1][-1])
    
    return results

p = Pool(mp.cpu_count())

# Pool.map() takes a function and an iterable like a list or generator
results_list = p.map(get_customer_data, [data_tuple for data_tuple in customer_data.itertuples(name=None, index=False)] )
# results is a list of dict()
results_dict = {k:v for x in results_list for k,v in x.items()}
p.close()

About the Author

Question Info

Publish Time
2023-01-25 16:34
Update Time
2023-01-25 16:34