How to get corresponding numpy array values after performing calculation on another related array

Question

I have a function that returns two NumPy arrays (width and height) like so:

width, height = calc_heigh_width(data)

width
>>> array([390, 20, 65, 1000])

height
>>> array([2, 7, 3, 1])

Imagine the widths as being points on an x-axis, so they go from 0 to 1000 in this example, similarly for the height.

I am only interested in places where the current width is bigger than the previous width by 600 and less than the previous width by 800. So I wrote this to solve that:

import numpy as np

sorted_width_array = np.sort(width)

width_calc_list = []
for i in range(len(sorted_width_array) - 1):
    if ((sorted_width_array[i+1] - sorted_width_array[i]) > 600) and ((sorted_width_array[i+1] - sorted_width_array[i]) < 800):
        width_calc_list.append(sorted_width_array[i])
        width_calc_list.append(sorted_width_array[i+1])

This returns

width_calc_list
>>> [390, 1000]

However, I would also like to get the height values that correspond to these width values but I just haven't been able to solve it. Any help is much appreciated.

So basically, I'd like to get something like this

height_calc_list
>>> [2, 1]

PS: I know my current code to find the width works but if there's any way to improve it I'm all ears. I was playing around with NumPy's diff function and I'm optimistic that will work but haven't been able to get it to work.

Answer 1

you can just zip the values and use the same algorithm you have written. i did not check the correctness of your solution, just changed it so that it returns the format you want. check it out

width = [390, 20, 65, 1000]
height = [2, 7, 3, 1]

cord = sorted(list(zip(width, height)), key=lambda tup: tup[0])
print(cord)

width_calc_list = []
for i in range(len(cord) - 1):
    if ((cord[i+1][0] - cord[i][0]) > 600) and ((cord[i+1][0] - cord[i][0]) < 800):
        width_calc_list.append(cord[i])
        width_calc_list.append(cord[i+1])
        
print(width_calc_list)
#[(390, 2), (1000, 1)]