Блог

У меня есть фрейм данных в Pandas, содержащий названия планет и информацию о каждой планете. Названия планет хранятся под столбцом с заголовком «планета». Я хочу выделить строки конкретных планет, Земли, Сатурна и Марса. По какой-то причине, когда я запускаю этот скрипт и экспортирую в виде файла Excel, выделяется только первая планета в моем списке highlight_planets. В этом случае выделяется только Земля (Сатурн и Марс не выделяются). Если я переместу Saturn в первую позицию, а Earth во вторую позицию, во фрейме данных будут выделены только строки Saturn. Как мне выделить все строки, в которых мой столбец планета содержит Землю, Сатурн или Марс?

Спасибо.

 def highlight_sentiment(row):

    highlight_planets = ['Earth', 'Saturn', 'Mars']

    for m in highlight_planets:


        if row['planet'] == m:
           return ['background-color: yellow'] * len(row)
        else:
           return ['background-color: white'] * len(row)

df.style.apply(highlight_sentiment, axis=1)

df = df.style.apply(highlight_sentiment, axis=1)
df.to_excel("df_test.xlsx", index = False)
 

Более стандартный подход к этой проблеме заключается в создании фрейма данных стилей, который можно создавать динамически на основе расположения планет в списке, используя loc Series.isin :

 def highlight_sentiment(df_, highlight_planets):
    # Build DataFrame With Default Style
    styles_df = pd.DataFrame('background-color: white',
                             index=df_.index,
                             columns=df_.columns)
    # Update Styles based on condition
    styles_df.loc[
        df_['planet'].isin(highlight_planets), :
    ] = 'background-color: yellow'
    return styles_df


# Pass highlight_planets as an argument to allow more flexible styles
styler = df.style.apply(highlight_sentiment,
                        highlight_planets=['Earth', 'Saturn', 'Mars'],
                        axis=None)
# Use styler to write to excel
styler.to_excel("df_test.xlsx", index=False)
 

Примеры данных и импорта:

 import pandas as pd

df = pd.DataFrame({
    'planet': ['Mercury', 'Venus', 'Earth',
               'Mars', 'Jupiter', 'Saturn'],
    'data': [1, 2, 3, 4, 5, 6]
})
 

В вашем for цикле вы return сразу же, если планета не является первым значением в вашем списке.

Переместите последний return оператор за пределы for цикла, чтобы он возвращал белый фон только в том случае, если планета не соответствует ни одному значению в вашем списке.

 def highlight_sentiment(row):
    highlight_planets = ['Earth', 'Saturn', 'Mars']
    for m in highlight_planets:
        if row['planet'] == m:
           return ['background-color: yellow'] * len(row)
    return ['background-color: white'] * len(row)
 

Кроме того, вы могли for бы вообще избежать цикла и просто использовать in :

 def highlight_sentiment(row):
    highlight_planets = ['Earth', 'Saturn', 'Mars']
    if row['planet'] in highlight_planets:
        return ['background-color: yellow'] * len(row)
    return ['background-color: white'] * len(row)