Finding tree index from list of positions

[hyanwong]

Just to remind myself, I think this does it: worth noting because it's easy to forget the side="right" and -1

import numpy as np
list_of_positions = [1, 2, 10, 2, 4]
tree_index = np.searchsorted(ts.breakpoints(as_array=True), list_of_positions), side='right') - 1

[hanbin973]

Funny that I spent an hour implementing this by hand :(

@numba.njit
def _edges_first_tree(edges_left, breakpoints):
    edges_first_tree = np.empty(edges_left.shape[0])
    for e, left in enumerate(edges_left):
        # binary search
        tree_left, tree_right = 0, breakpoints.shape[0]-1
        while True:
            tree_now = int((tree_left+tree_right)/2)
            if breakpoints[tree_now] < left:
                tree_left = tree_now
            elif breakpoints[tree_now] > left:
                tree_right = tree_now 
            else:
                edges_first_tree[e] = tree_now
                break
    return edges_first_tree

[hyanwong]

I guess it's useful to have 2 implementations to check they give the same result. Is the numba version faster or slower, by the way?

[hanbin973]

They provide the same result. The numpy version is faster in my test case.

[jeromekelleher]

It'll be hard to beat numpy's implementation of binary search, but impressive that your numba version gets so close @hanbin973 !