How to compare two columns in a dataframe with a function and for loop?

Question

I have two columns floors and floor total

df = pd.DataFrame({"floor": [1,2,30], "floors_total": [1, 50, 30]})

I want to write a function such as if in a row both values in floors total and floor are equal then return "last floor" and then if floor = 1 then return "first floor" and if other return "other" I want to apply this to df

I am trying to write a for loop which is :

def new(df):
  a = "floor"
  b = "floors_total"
  for a,b in df.iteritems():
    if a == b:
      return "last floor"
    elif a == 1:
      return "first floor"
    else:
      return "other"
df["floor_postiton"] = df.apply(new)

I am getting all rows in the new column as NaN

Please follow the [guidelines](https://stackoverflow.com/help/how-to-ask) to ask a proper question on the forum — Yayati Sule, Jun 16 '20 at 07:45
What is the reason to use a (slow) iteration instead of the vectorized pandas or numpy methods? — Serge Ballesta, Jun 16 '20 at 08:12
No reason. I have not yet learnt vectorized pandas or numpy methods. Is there faster way with numpy code? — Neekunj, Jun 16 '20 at 08:13

score 1 · Accepted Answer · answered Jun 16 '20 at 08:39

1

You can avoid the (slow) iteration over the dataframe with np.where:

df['floor_position'] = np.where(df['floor'] == df['floors_total'],
                'last floor',
                np.where(df['floor'] == 1, 'first floor',
                     'other'))

answered Jun 16 '20 at 08:39

Serge Ballesta

143,923
11
122
252

How to compare two columns in a dataframe with a function and for loop?

1 Answers1