给定值列表迭代过滤 Pandas 数据框的最有效方法

Question

我希望有人能给我指出正确的方向...

我有一个包含市场数据的数据集，其中包含以下列：面积、每晚价格和卧室数量。

我想创建另一个数据集，显示每个区域在给定列表中的每个价格点有多少卧室可用。

我目前使用两个 for 循环：一个循环遍历价格，另一个循环遍历区域（循环中的循环）。然后过滤市场数据并汇总卧室列。这是一个非常缓慢的过程 - 特别是当我的价目表有 1000 个条目并且我有几十个区域时。

我怎样才能加快这个过程？下面附上示例代码。

import random
import pandas as pd

name_choices = ['North', 'South', 'East', 'West']
bedroom_choices = [1,2,3,4,5]
price_choices = list(range(5, 300))

name_list = []
bedrooms_list = []
price_list = []

for i in range(100):
    
    name_list.append(random.choice(name_choices))
    bedrooms_list.append(random.choice(bedroom_choices))
    price_list.append(random.choice(price_choices))

market_data_ex  = pd.DataFrame(data = {'Area' : name_list, 'Bedrooms' : bedrooms_list, 'Price': price_list})

empty_area = []
empty_price = []
empty_bedrooms = []

for area in name_choices:
    
    for price in price_choices:
        
        bedrooms_available = market_data_ex[(market_data_ex['Area'] == area) & (market_data_ex['Price'] <= price)]['Bedrooms'].sum()
        
        empty_area.append(area)
        empty_price.append(price)
        empty_bedrooms.append(bedrooms_available)
        
pd.DataFrame(data = {'Area' : empty_area, 'Price' : empty_price, 'Bedrooms' : empty_bedrooms})

非常感谢!!!

Answer 1

如果我没理解错的话，你可以做.pivot_table你的数据，然后做.cumsum() row-wise (axis=1):

x = (
    market_data_ex.pivot_table(
        index="Area",
        columns="Price",
        values="Bedrooms",
        aggfunc="sum",
        fill_value=0,
    )
    .cumsum(axis=1)
    .stack()
    .reset_index()
    .rename(columns={0: "Bedrooms"})
)
print(x)

给定值列表迭代过滤 Pandas 数据框的最有效方法

Most Efficient Way to iteratively filter a Pandas dataframe given a list of values

python

loops

vectorization

python-3.x