CODEWITHSUNDEEP
pythonpandas
Don Giulio
Don Giulio

Reputation: 3284

pandas update dataframe only when all similar values are valid

I've written this simplified example to explain what I'm trying to achieve: 

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)

The distance dataframe has several contingent rows having the same distance and letter columns. initially they are all enabled == False

I need to set some of them enabled == True based on their row index and a range value, so that all rows with a range distance from the one with index row will be enabled (and this I managed to get in my enable_rows function).

Additionally I need that if one distance value wouldn't have all its rows enabled then none should be enabled.

both examples in the code above would have some of the distance == 100.0 rows still not enabled, so none of the 100.0 should be enabled.

They expect a resulting dataframe as : 

   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

but the actual output of the program is: 

   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

how could I update enable_rows to obtain that?

Upvotes: 1

Views: 37

Answers (1)

b-fg
b-fg

Reputation: 4137

You just need to groupby 'distances' and transform the result if all the enabled values are not True. You can do this with:

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

Which you can use here

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)

Upvotes: 1

Related Questions