Панды обновляют фрейм данных только тогда, когда все подобные значения действительны

Question

Я написал этот упрощенный пример, чтобы объяснить, чего я пытаюсь достичь:

import pandas as pd
import pytest

def enable_rows(df, row, myrange):
    # Need to modify this
    df.loc[row + myrange:, 'enabled'] = True
    df.loc[:row - myrange, 'enabled'] = True

def starting_df():
    # just re-creates the initial dataframe to check on values
    distance = {1: (100.0, 'a', False),
                2: (100.0, 'a', False),
                3: (100.0, 'a', False),
                4: (700.0, 'b', False),
                5: (700.0, 'b', False),
                6: (900.0, 'c', False)}

    return pd.DataFrame(data=list(distance.values()), index=list(
        distance.keys()), columns=['distance', 'letter', 'enabled'])

def test_enable(center_row, myrange):
    # convenience function to eye-candy the executions.
    df = starting_df()
    enable_rows(df, center_row, myrange)
    print(df)

    # assertions
    enabled = df.loc[df.enabled]
    if not ((len(enabled) == 3) and
            (len(enabled.loc[df.distance == 100.0]) == 0) and
            (len(enabled.loc[df.distance > 100.0]) == 3)):
        print("wrong result")

test_enable(1, 2)
test_enable(2, 1)

В кадре данных расстояния есть несколько условных строк, имеющих одинаковые столбцы distance и letter.изначально они все enabled == False

Мне нужно установить некоторые из них enabled == True на основе их индекса row и значения range, чтобы все строки на расстоянии range от одногос индексом row будет включено (и это мне удалось получить в моей функции enable_rows).

Кроме того, мне нужно, чтобы, если для одного значения distance не было включено все его строки, ни одна из них не должна быть включена.

в обоих примерах в приведенном выше коде есть некоторые из distance == 100.0строки все еще не включены, поэтому ни один из 100.0 не должен быть включен.

Они ожидают результирующий кадр данных как:

   distance letter  enabled
1     100.0      a    False
2     100.0      a    False
3     100.0      a    False
4     700.0      b     True
5     700.0      b     True
6     900.0      c     True

, но фактический результат программы:

   distance letter  enabled
1     100.0      a    False
2     100.0      a    False
3     100.0      a     True
4     700.0      b     True
5     700.0      b     True
6     900.0      c     True
wrong result
   distance letter  enabled
1     100.0      a     True
2     100.0      a    False
3     100.0      a     True
4     700.0      b     True
5     700.0      b     True
6     900.0      c     True
wrong result

как я могу обновить enable_rows, чтобы получить это

Answer 1

Вам просто нужно groupby 'distances' и transform результат, если все значения enabled не True.Вы можете сделать это с помощью:

df['enabled'] = df.groupby('distance')['enabled'].transform(lambda x: all(x)==True)

, которые вы можете использовать здесь

def enable_rows(df, row, myrange):
    # Need to modify this
    df.loc[row + myrange:, 'enabled'] = True
    df.loc[:row - myrange, 'enabled'] = True
    df['enabled'] = df.groupby('distance')['enabled'].transform(lambda x: all(x)==True)