如何基于MySQL结果在Python中按条件输出各机构Top2主分类及子分类

Alright, let's walk through how to solve this problem step by step. We'll use pandas for this—it's ideal for grouped ranking tasks like this. First, let's clarify the ranking logic we need to follow:

1. For each org_id, calculate the total number for each main_cat, then pick the top 2 main_cat with the highest total values.
2. For each of those top 2 main_cat, select the top 2 sub_cat based on their individual number values.
3. This ensures each org_id ends up with at most 4 records (2 main categories × 2 sub categories).

Step 1: Prep the Raw Data #

First, we need to convert your pipe-separated raw data into a structured pandas DataFrame. Here's how to do that:

import pandas as pd

# Your raw data, cleaned up for easier processing
raw_data = """main_cat| sub_cat | number | org_id
Career | school | 5 | A
Career | college | 3 | A
Career | higher | 4 | A
Job | Blr | 6 | A
Job | Hyd | 11 | A
Job | Chennai | 12 | A
Career | school | 15 | B
Career | college | 30 | B
Career | higher | 5 | B
Job | Blr | 5 | B
Career | college | 8 | C
Job | Chennai | 4 | C"""

# Split into lines and clean up each entry
lines = [line.strip() for line in raw_data.split('\n') if line.strip()]
column_names = [col.strip() for col in lines[0].split('|')]
data_rows = []
for line in lines[1:]:
    # Split each line by | and remove extra spaces
    row_parts = [part.strip() for part in line.split('|')]
    data_rows.append(row_parts)

# Create DataFrame and convert the 'number' column to integers (since we'll do math on it)
df = pd.DataFrame(data_rows, columns=column_names)
df['number'] = df['number'].astype(int)

Step 2: Rank Main Categories per Organization #

Next, we'll calculate the total number for each main category within each organization, then rank them to pick the top 2:

# Calculate total number for each org_id + main_cat combination
main_cat_totals = df.groupby(['org_id', 'main_cat'])['number'].sum().reset_index(name='total_number')

# Assign a rank to each main_cat within its org_id (highest total gets rank 1)
main_cat_totals['main_cat_rank'] = main_cat_totals.groupby('org_id')['total_number'].rank(ascending=False, method='dense')

# Filter to keep only the top 2 main_cats per org_id
top_main_cats = main_cat_totals[main_cat_totals['main_cat_rank'] <= 2][['org_id', 'main_cat']]

Step 3: Rank Sub Categories for Top Main Categories #

Now we'll filter the original data to only include our top main categories, then rank the sub categories within each of those:

# Merge the top main_cats list back with the original data to filter out unwanted main_cats
filtered_data = pd.merge(df, top_main_cats, on=['org_id', 'main_cat'])

# Rank each sub_cat within its org_id + main_cat group (highest number gets rank 1)
filtered_data['sub_cat_rank'] = filtered_data.groupby(['org_id', 'main_cat'])['number'].rank(ascending=False, method='dense')

# Keep only the top 2 sub_cats per org_id + main_cat
final_result = filtered_data[filtered_data['sub_cat_rank'] <= 2][['org_id', 'main_cat', 'sub_cat', 'number']]

Step 4: Print the Final Result #

Finally, let's print the result in a sorted, readable format:

# Sort by org_id, then main_cat, then sub_cat rank for clarity
sorted_result = final_result.sort_values(by=['org_id', 'main_cat', 'sub_cat_rank'])
print("Final Output:")
print(sorted_result)

What the Output Looks Like #

When you run this code, you'll get:

org_id main_cat   sub_cat  number
5      A      Job   Chennai      12
4      A      Job       Hyd      11
0      A   Career    school       5
2      A   Career    higher       4
7      B   Career   college      30
6      B   Career    school      15
9      B      Job       Blr       5
10     C   Career   college       8
11     C      Job   Chennai       4

Let's break this down:

• Org A: Top main_cats are Job (total 29) and Career (total 12). Job's top sub_cats are Chennai (12) and Hyd (11); Career's are school (5) and higher (4).
• Org B: Top main_cats are Career (total 50) and Job (total 5). Career's top sub_cats are college (30) and school (15); Job only has one sub_cat, so it's included.
• Org C: There are only 2 main_cats, so both make the top 2. Each has one sub_cat, so both are kept.

More Concise Alternative #

If you want to skip the merge step, you can use pandas' transform method to calculate ranks directly on the original DataFrame:

# Calculate total number per main_cat (within each org) and assign rank
df['main_cat_total'] = df.groupby(['org_id', 'main_cat'])['number'].transform('sum')
df['main_cat_rank'] = df.groupby('org_id')['main_cat_total'].rank(ascending=False, method='dense')

# Filter to top 2 main_cats
df_top_main = df[df['main_cat_rank'] <= 2]

# Rank sub_cats and filter to top 2
df_top_main['sub_cat_rank'] = df_top_main.groupby(['org_id', 'main_cat'])['number'].rank(ascending=False, method='dense')
final_result = df_top_main[df_top_main['sub_cat_rank'] <= 2][['org_id', 'main_cat', 'sub_cat', 'number']]

# Print sorted result
print(final_result.sort_values(by=['org_id', 'main_cat', 'sub_cat_rank']))

This gives exactly the same result, just in a more streamlined way.

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