如何基于Pandas DataFrame指定列的唯一值为整行添加背景色？（续问）

Hey there! Let's fix that styling issue you're having with Pandas DataFrames. The error you're seeing happens because when using apply(axis=1), your function needs to return a list/array that matches the number of columns in each row—not just a single string. Here's how to get entire rows highlighted based on unique "Country" values properly:

The Root of the Error #

Your original test_check function returns a single style string per row, but Pandas expects a sequence (one value for each column in the row). That's why you get the shape mismatch: it expects a (5,3) output (5 rows × 3 columns) but only gets a (5,) output.

Optimal Solution #

We'll create a color map for unique countries, then define a function that returns a style string for every column in a matched row. Here's the clean, working code:

import pandas as pd
import matplotlib

# Your sample DataFrame
d1 = pd.DataFrame({
    "Country": ['xx','xx','xy','xz','xz'], 
    "year": ['y1','y2','y1','y1','y2'], 
    "population": [100,200,120,140,190]
})

# Map unique Country values to distinct colors from matplotlib
unique_countries = d1['Country'].unique()
# Slice color list to match the number of unique countries (avoid unused colors)
color_list = list(matplotlib.colors.cnames.values())[:len(unique_countries)]
country_color_map = dict(zip(unique_countries, color_list))

# Function to generate styling for an entire row
def highlight_row_by_country(row):
    # Grab the color for the current row's Country
    row_color = country_color_map[row['Country']]
    # Return a style string for every column in the row
    return [f'background-color: {row_color}'] * len(row)

# Apply the styling to the DataFrame
styled_df = d1.style.apply(highlight_row_by_country, axis=1)
styled_df

Alternative Efficient Approach (For Larger DataFrames) #

If you're working with bigger datasets, creating a style DataFrame upfront and applying it all at once is more efficient:

# Create a DataFrame where each cell has the style string for its row's Country
style_template = d1['Country'].map(country_color_map).apply(
    lambda color: [f'background-color: {color}'] * d1.shape[1]
)

# Apply the pre-built style DataFrame
styled_df = d1.style.apply(lambda _: style_template, axis=None)
styled_df

Both methods will give you the desired result: every row gets a unique background color based on its "Country" value, with all columns in the row highlighted.

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