Find the next higher or lower number in panda

Viewed 147

I want to add a new column to my DataFrame with the next higher or lower number to y from an external array with allowed_numbers_y.

I have:

allowed_numbers_y = [4,8,15,47]

df = pd.DataFrame({'x': [1,1,1,2,2,3,3,3,4,4,4],
                   'y': [5,12,13,14,15,16,17,25,21,10,9]})

df output looks like this:

    x   y
0   1  5
1   1  12
2   1  13
3   2  14
4   2  15
5   3  16
6   3  17
7   3  25
8   4  21
9   4  10
10  4   9

I would like to have this:

    x   y    next_up_y  next_down_y
0   1  5        8         4
1   1  12       15        8
2   1  13       15        8
3   2  14       15        8
4   2  15       47        15
5   3  16       47        15
6   3  17       47        15
7   3  25       47        15
8   4  21       47        15
9   4  10       15        8
10  4   9       15        8

I have this, but since my dataframe is big I think there must be a better thing to do instead of looping over every row with this function.

def next_higher(my_list, my_number):  # finds the next higher number in list

    jojo_list = list(my_list)
    if not (len(list(filter(lambda x: x == my_number, my_list))) > 0):  # Number exists
        jojo_list.append(int(my_number))
    jojo_list.sort()
    jojo_list.index(int(my_number))
    if my_number >= my_list[-1]:  # my_number is the last in list
        return my_list[-1]
    else:
        return jojo_list[jojo_list.index(my_number) + 1]
2 Answers

If you convert the allowed list into a Series, you can use idxmin() with apply().

For next_up_y:

  • allowed[allowed > y] - get values greater than y
  • .sub(y).abs() - take the absolute value of the difference with y
  • .idxmin() - return the index of the min value

For next_down_y, do the same thing but with values less than y (allowed[allowed < y]).

allowed = pd.Series(allowed_numbers_y)

df['next_up_y'] = df.y.apply(
    lambda y: allowed[allowed[allowed > y].sub(y).abs().idxmin()])
df['next_down_y'] = df.y.apply(
    lambda y: allowed[allowed[allowed < y].sub(y).abs().idxmin()])
    x   y  next_up_y  next_down_y
0   1   5          8            4
1   1  12         15            8
2   1  13         15            8
3   2  14         15            8
4   2  15         47            8
5   3  16         47           15
6   3  17         47           15
7   3  25         47           15
8   4  21         47           15
9   4  10         15            8
10  4   9         15            8

Note: If speed is critical, check out @Durtal's numpy solution, which is less intuitive but orders-of-magnitude faster:

method %timeit
pandas 11 ms ± 1.46 ms per loop (7 runs, 100 loops each)
numpy 16.2 µs ± 1.27 µs per loop (7 runs, 10000 loops each)

If speed is of importance, you have to use numpy. It is more cryptic than tdy's proposal, but I bet it is faster. The following calculation assumes that min(y_allowed) <= y <= max(y_allowed) and that y_allowed is sorted. Otherwise, some simple adjustments are needed

import numpy as np

allowed_numbers_y = [4,8,15,47]

y = np.array([5,12,13,14,15,16,17,25,21,10,9])

# Calculate diffence between y and all allowed_numbers_y
allowed_numbers_extended = np.array([allowed_numbers_y,]*len(y))
matrix_difference = allowed_numbers_extended - y[:,np.newaxis]

# find index of first element greater zero in every row
matrix_greater_zero = np.less(0, matrix_difference)
idx_higher = matrix_greater_zero.argmax(axis=1) 
# index of last element less than zero is one less
idx_lower = idx_higher - 1

next_up_y = np.array(allowed_numbers_y)[idx_higher]
next_down_y = np.array(allowed_numbers_y)[idx_lower]
Related