Add a column based on a condition in Pandas

I have the following df.

data = {
    'policyId': ['X','X','X','X', 'Y','Y','Y','Y', 'Z', 'Z','Z','Z'],
    'element_id': [242, 243, 241, 257, 242, 243, 241, 257, 242, 243, 241, 257],
    'element_selected': [True, True, True, True, False, True, True, True, True, True, True, False]
}
df = pd.DataFrame(data)

All policyIds have the same element_id. What changes is that the element_id can be select or not (element_selected assigns that). I want to add a contract_type column to the df. If element_id 242 has been selected the contract_type should be "CONTENT". If element_id 257 has been selected then the contract_type should "PRIVATE LIABILITY". If element_id 242 and 257 have both been selected then the contract_type should "CONTENT + PRIVATE LIABILITY".

The problem is that each policyId should have only one value per contract_type and I do not know how to implement this logic

Expected output

policyId	element_id	element_selected	contract_type
X	               242	True	            CONTENT + PRIVATE LIABILITY
X	               243	True	            CONTENT + PRIVATE LIABILITY
X	               241	True	            CONTENT + PRIVATE LIABILITY
X	               257	True	            CONTENT + PRIVATE LIABILITY
Y	               242	False	            PRIVATE LIABILITY
Y	               243	True	            PRIVATE LIABILITY
Y	               241	True	            PRIVATE LIABILITY
Y	               257	True	            PRIVATE LIABILITY
Z	               242	True	            CONTENT
Z	               243	True	            CONTENT
Z	               241	True	            CONTENT
Z	               257	False	            CONTENT

Solution

Pivot the dataframe

s = df.pivot(index='policyId', columns='element_id', values='element_selected')

# element_id   241    242   243    257
# policyId                            
# X           True   True  True   True
# Y           True  False  True   True
# Z           True   True  True  False

Use np.select to select the choices based on predefined conditions

s['contract_type'] = np.select(
    
展开收缩
].all(axis=1), s[242], s[257]],
    ["CONTENT + PRIVATE LIABILITY", "CONTENT", "PRIVATE LIABILITY"]
)

# element_id   241    242   243    257                contract_type
# policyId                                                         
# X           True   True  True   True  CONTENT + PRIVATE LIABILITY
# Y           True  False  True   True            PRIVATE LIABILITY
# Z           True   True  True  False                      CONTENT

Map the selected choices back to dataframe

df['contract_type'] = df['policyId'].map(s['contract_type'])

#    policyId  element_id  element_selected                contract_type
# 0         X         242              True  CONTENT + PRIVATE LIABILITY
# 1         X         243              True  CONTENT + PRIVATE LIABILITY
# 2         X         241              True  CONTENT + PRIVATE LIABILITY
# 3         X         257              True  CONTENT + PRIVATE LIABILITY
# 4         Y         242             False            PRIVATE LIABILITY
# 5         Y         243              True            PRIVATE LIABILITY
# 6         Y         241              True            PRIVATE LIABILITY
# 7         Y         257              True            PRIVATE LIABILITY
# 8         Z         242              True                      CONTENT
# 9         Z         243              True                      CONTENT
# 10        Z         241              True                      CONTENT
# 11        Z         257             False                      CONTENT

you could use a groupy "policyId" and then implement a mapping function that does the logic of assigning the contract_type:

def get_contract_type(group):
    types = []
    # get selected ids
    ids = list(group[group["element_selected"]]["element_id"])

    # add the types based on the selected ids
    if 242 in ids:
        types.append("CONTENT")

    if 257 in ids:
       types.append("PRIVATE LIABILITY")


    # add the new contract_type column
    group["contract_type"] = " + ".join(types)
    return group

df.groupby("policyId").apply(get_contract_type)

First, let's create a temporary column that contains zeros:

df["contract_type"] = 0

The purpose of initializing this column to zeros will become clear later in this answer.

Now, group by policyId:

policy_groups = df.groupby("policyId")

Next, for each group, filter out the rows where element_selected is False, and see if 242 and 257 exist in the element_id column. If 242 exists, let's add 1 to contract_type for all rows in the group. If 257 exists, let's add 2.

for policy_id, group in policy_groups:
    selected_ids = group.loc[group["element_selected"], "element_id"]
    if (selected_ids == 242).any():
        df.loc[group.index, "contract_type"] += 1
    if (selected_ids == 257).any():
        df.loc[group.index, "contract_type"] += 2 

At this stage, df looks like:

   policyId  element_id  element_selected  contract_type
0         X         242              True              3
1         X         243              True              3
2         X         241              True              3
3         X         257              True              3
4         Y         242             False              2
5         Y         243              True              2
6         Y         241              True              2
7         Y         257              True              2
8         Z         242              True              1
9         Z         243              True              1
10        Z         241              True              1
11        Z         257             False              1

Now that this is complete for all groups, we can map the integers in the contract_type column to the required strings:

df["contract_type"] = df["contract_type"].map({0: "", 1: "CONTENT", 2: "PRIVATE LIABILITY", 3: "CONTENT + PRIVATE LIABILITY"})

And we have the required dataframe:

   policyId  element_id  element_selected                contract_type
0         X         242              True  CONTENT + PRIVATE LIABILITY
1         X         243              True  CONTENT + PRIVATE LIABILITY
2         X         241              True  CONTENT + PRIVATE LIABILITY
3         X         257              True  CONTENT + PRIVATE LIABILITY
4         Y         242             False            PRIVATE LIABILITY
5         Y         243              True            PRIVATE LIABILITY
6         Y         241              True            PRIVATE LIABILITY
7         Y         257              True            PRIVATE LIABILITY
8         Z         242              True                      CONTENT
9         Z         243              True                      CONTENT
10        Z         241              True                      CONTENT
11        Z         257             False                      CONTENT

Try this:

(df.assign(contract_type = df['policyId']
           .map(df['element_id'].map({242:'0',257:'1'})
                .where(df['element_selected'])
                .str.get_dummies()
                .groupby(df['policyId']).any()
                .mul([1,2])
                .sum(axis=1)
                .map(dict(enumerate(['NEITHER','CONTENT','PRIVATE LIABILITY','CONTENT + PRIVATE LIABILITY']))))))

or

df.assign(
    contract_type = (df['policyId'].map(
        df['element_id']
        .sort_values()
        .map(
            {242:'CONTENT',257:'PRIVATE LIABILITY'})
            .where(df['element_selected'] & ~df.duplicated(subset = ['policyId','element_id']))
            .dropna()
            .groupby(df['policyId'])
            .agg(' + '.join))))

Output:

   policyId  element_id  element_selected                contract_type
0         X         242              True  CONTENT + PRIVATE LIABILITY
1         X         243              True  CONTENT + PRIVATE LIABILITY
2         X         241              True  CONTENT + PRIVATE LIABILITY
3         X         257              True  CONTENT + PRIVATE LIABILITY
4         Y         242             False            PRIVATE LIABILITY
5         Y         243              True            PRIVATE LIABILITY
6         Y         241              True            PRIVATE LIABILITY
7         Y         257              True            PRIVATE LIABILITY
8         Z         242              True                      CONTENT
9         Z         243              True                      CONTENT
10        Z         241              True                      CONTENT
11        Z         257             False                      CONTENT

Let's do it with CategoricalDtype

ContractTypes = {
    242: 'Content',
    257: 'Private liability'
}

contract_type = (
    df.set_index('policyId')
    .product('columns')
    .astype(pd.CategoricalDtype(ContractTypes))
    .cat.rename_categories(ContractTypes.values())
    .dropna()
    .groupby(level=0)
    .agg(' + '.join)
    .rename('contract_type')
)

df = df.set_index('policyId').join(contract_type).reset_index()

First we create a dictionary of contract type codes as keys and associated descriptions. With time we may want to add a new contract type here without mitigating with the rest of the code.

Next we create a series contract_type:

• put policyId as an index for further grouping by index;
• product along the rows returns zero in place of false codes;
• transforming the result to categorical ContractTypes will replace unused contract types with NaN, which in turn we can safely discard, since the data we have is tied to the policyId index;
• next we rename categories with their associated descriptions;
• now we group the rest of the data by index (i.e. level=0) and join them together with ' + ' in between;
• finally we give the sequence a name which will become a name of a column after further joining with df.

Let's do it by Pivoting

ContractTypes = pd.Series({
    242: 'Content',
    257: 'Private liability'
})

selected = df.pivot(
    index='policyId',
    columns='element_id',
    values='element_selected'
).apply(
    lambda s: ' + '.join(ContractTypes
展开收缩
), 
    axis='columns'
)

df['contract_type'] = df['policyId'].map(selected)

This one is an upgraded version of the Shubham Sharma's answer.
The only change is made in the middle, where each row of the pivoted data is used as an index in the ContractTypes series, and obtained descriptions are joined in a string.