Generalize a function in python

Question

I need to calculate some probabilities given certain conditions, so I'm using a function to get the rows that contain givens values, for example:

df:
col1 col2 col3
 A    B    C
 H    B    C
 A    B
 A    H    C

This is the function

def existence(x):
    return df[df.isin([x]).any(1)]

So if I do:

in:
existence('A')

out:
col1 col2 col3
 A    B    C
 A    B
 A    H    C

I need to generalize the function, so that I can give it more than one parameter and do the following:

existence(x, y):
    return df[df.isin([x]).any(1) & df.isin([y]).any(1)]

or generalized

existence(x1, x2,..., xn):
    return df[df.isin([x1]).any(1) & df.isin([x2]).any(1) & ... & df.isin([xn]).any(1)]

I think args can't help me since I can't merge operations with the operator &

thank you in advance estimates

def existence(x): return df[df.isin(x).any(1)], then existence(['A']) or existence(['A','B']) — BENY
– BENY, Commented May 16, 2019 at 20:37

user3483203 · Accepted Answer · 2019-05-16 20:28:05Z

8

Your function will already work with a variable number of arguments. Check out the description of pd.DataFrame.isin:

DataFrame.isin(values) Whether each element in the DataFrame is contained in values.

This works fine with more than one parameter, you just need to change how you pass parameters into your function. This is what *args is for in Python. It allows you to pass in a variable number of arguments and gives you back a tuple.

I also updated your function to take in the DataFrame to apply the mask on, because it isn't good pratice to rely on global variable names.

def existence(df, *args):
    return df[df.isin(args).any(1)]

In [13]: existence(df, 'A')
Out[13]:
  col1 col2  col3
0    A    B     C
2    A    B  None
3    A    H     C

In [14]: existence(df, 'C', 'H')
Out[14]:
  col1 col2 col3
0    A    B    C
1    H    B    C
3    A    H    C

edited May 16, 2019 at 20:28

answered May 16, 2019 at 20:22

user3483203

51.3k10 gold badges72 silver badges104 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Generalize a function in python

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related