python - 仅当所有相似值都有效时，pandas 才会更新数据框

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)

距离数据帧有几个具有相同距离和字母列的偶然行。最初它们都是enabled == False

我需要根据它们的row索引和range值设置其中一些enabled == True，以便所有行将启用距索引 row 的 range 距离(我设法在我的 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'并转换结果正确。您可以通过以下方式执行此操作:

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)