Python Dataframe yes/no checker

Viewed 113

I would like to make a table that evaluates whether a user is in a group or not. How can I get my dictionary sorted like in the example I have down below? I would like the index and columns populated automatically by the key and value.

d = {
    'user1': ['group1', 'group2', 'group3'],
    'user2': ['group1', 'group2'],
    'user3': ['group2']}

df = pd.DataFrame.from_dict(d,orient='index')

print(df)

Current result

    user1   user2   user3
0  group1  group1  group2
1  group2  group2     NaN
2  group3     NaN     NaN

Desired result - I would like the rows based on the key and the columns based on the values.

        group1  group2  group3
user1   Y       Y       Y
user2   Y       Y       N
user3   N       Y       N
4 Answers

Try:

d = {
    "user1": ["group1", "group2", "group3"],
    "user2": ["group1", "group2"],
    "user3": ["group2"],
}

df = pd.DataFrame.from_dict(d, orient="index")

x = df.stack().droplevel(level=1)
x = pd.crosstab(x.index, x).replace({1: "Y", 0: "N"})
x.index.name, x.columns.name = None, None
print(x)

Prints:

      group1 group2 group3
user1      Y      Y      Y
user2      Y      Y      N
user3      N      Y      N

Convert your dict to a list of tuples with a comprehension:

>>> pd.DataFrame([(u, g, 'Y') for u, lg in d.items() for g in lg]) \
      .pivot(0, 1, 2).fillna('N').rename_axis(index=None, columns=None)

      group1 group2 group3
user1      Y      Y      Y
user2      Y      Y      N
user3      N      Y      N

We can stack to convert to long form, use str.get_dummies to one-hot encode each user, then groupby max to resolve back to a single row per user. Lastly, replace the binary data with Y and N via replace:

out_df = (
    df.stack().str.get_dummies()
        .groupby(level=0).max()
        .replace({1: 'Y', 0: 'N'})
)

out_df:

      group1 group2 group3
user1      Y      Y      Y
user2      Y      Y      N
user3      N      Y      N

You can melt the data frame to disaggragate it, drop NA values, then assign a dummy placeholder value. From there you can pivot the data to your desired output.

df2 = df.melt().dropna()
df2['dummy'] = True
df2.pivot('variable', 'value').fillna(False).reset_index(drop=True)
# returns:
           z
value group1 group2 group3
0       True   True   True
1       True   True  False
2      False   True  False

Related