Finding elements in a pandas dataframe

Question

I have a pandas dataframe which looks like the following:

0 1
0 2
2 3
1 4

What I want to do is the following: if I get 2 as input my code is supposed to search for 2 in the dataframe and when it finds it returns the value of the other column. In the above example my code would return 0 and 3. I know that I can simply look at each row and check if any of the elements is equal to 2 but I was wondering if there is one-liner for such a problem.

UPDATE: None of the columns are index columns.

Thanks

Answer 1

>>> df = pd.DataFrame({'A': [0, 0, 2, 1], 'B': [1,2,3,4]})
>>> df
   A  B
0  0  1
1  0  2
2  2  3
3  1  4

The following pandas syntax is equivalent to the SQL SELECT B FROM df WHERE A = 2

>>> df[df['A'] == 2]['B']
2    3
Name: B, dtype: int64

There's also pandas.DataFrame.query:

>>> df.query('A == 2')['B']
2    3
Name: B, dtype: int64

Answer 2

You may need this:

n_input = 2

df[(df == n_input).any(1)].stack()[lambda x: x != n_input].unique()
# array([0, 3])