如何按索引列组合对DataFrame数据列进行求和？

Let's break down how to solve this problem step by step. First, let's recap your scenario: you have a DataFrame with road sections, speeds, and associated data columns, and you need to generate every possible combination of speeds across sections, then sum the corresponding data values for each combination.

Step 1: Set Up Your Example Data #

First, let's replicate your sample DataFrame to work with:

import pandas as pd

df = pd.DataFrame({
    'section': ['A', 'A', 'B', 'B'],
    'speed': [10, 20, 10, 20],
    'Data1': [1.5, 1.0, 2.5, 2.0],
    'Data2': [2.5, 2.0, 3.5, 3.0]
})

Step 2: Solution for 2 Sections (Your Example) #

If you're only working with 2 sections (like A and B), a straightforward approach is to split the data per section, then perform a cross join to get all speed combinations, followed by summing the data columns:

1. Split data by section and restructure:

   # Isolate each section's data, with speed as the index
   df_a = df[df['section'] == 'A'].set_index('speed')[['Data1', 'Data2']].rename(columns=lambda col: f"{col}_A")
   df_b = df[df['section'] == 'B'].set_index('speed')[['Data1', 'Data2']].rename(columns=lambda col: f"{col}_B")
   

2. Generate all speed combinations (cross join):
   Pandas doesn't have a dedicated cross join function, but we can use merge(how='cross') to get every possible pair of speeds between A and B:

   cross_combinations = df_a.reset_index().merge(df_b.reset_index(), how='cross')
   

3. Calculate summed data columns:
   Add the corresponding Data1 and Data2 values from each section, then clean up the output to match your desired format:

   # Compute sums
   cross_combinations['Data1'] = cross_combinations['Data1_A'] + cross_combinations['Data1_B']
   cross_combinations['Data2'] = cross_combinations['Data2_A'] + cross_combinations['Data2_B']
   
   # Format the final result
   final_result = cross_combinations[['speed_x', 'speed_y', 'Data1', 'Data2']].rename(columns={'speed_x': 'A', 'speed_y': 'B'})
   

Running this will give you exactly the output you're looking for:

A   B  Data1  Data2
0  10  10    4.0    6.0
1  10  20    3.5    5.5
2  20  10    3.5    5.5
3  20  20    3.0    5.0

Step 3: Generalized Solution for Up to 8 Sections #

Since you mentioned you might have up to 8 sections, here's a scalable approach that works regardless of how many sections you have:

from itertools import product

# Group data by section and create a mapping of speed to data values for each section
section_data_map = {}
for section, group in df.groupby('section'):
    # Convert each section's data to a dict: {speed: {'Data1': value, 'Data2': value}}
    section_data_map[section] = group.set_index('speed')[['Data1', 'Data2']].to_dict('index')

# Get list of sections and their available speeds
sections_list = list(section_data_map.keys())
speed_options = [list(section_data_map[sec].keys()) for sec in sections_list]

# Generate all possible speed combinations across sections
all_combinations = product(*speed_options)

# Calculate summed data for each combination
result_rows = []
for combo in all_combinations:
    # Initialize sums to 0
    data_sums = {'Data1': 0, 'Data2': 0}
    # Iterate through each section-speed pair in the combination
    for sec, speed in zip(sections_list, combo):
        data_sums['Data1'] += section_data_map[sec][speed]['Data1']
        data_sums['Data2'] += section_data_map[sec][speed]['Data2']
    # Combine section-speed values and sums into a single row
    row = dict(zip(sections_list, combo))
    row.update(data_sums)
    result_rows.append(row)

# Convert to final DataFrame
final_df = pd.DataFrame(result_rows)

This approach will automatically handle any number of sections (up to 8, as you noted) and any number of speed options per section.

内容的提问来源于stack exchange，提问作者Daniel F