Given a (flattened) list of lists, which element of which list does i correspond to?

index_of_which_bin.py

def index_of_which_bin(bin_sizes, absolute_index, verbose=False):
    """
    Given the absolute index, returns which bin it falls in and which element of that bin it corresponds to.
    """
    # Which class/bin does i fall into?
    accum = np.add.accumulate(bin_sizes)
    if verbose:
        print("accum =", accum)
    bin_index  = len(np.argwhere(accum <= absolute_index))
    if verbose:
        print("class_label =", bin_index)
    # Which element of the fallent class/bin does i correspond to?
    index_wrt_class = absolute_index - np.insert(accum, 0, 0)[bin_index]
    if verbose:
        print("index_wrt_class =", index_wrt_class)
        
    return bin_index, index_wrt_class

# Example
index_of_which_bin(bin_sizes = [10, 20, 30, 100], absolute_index = 15, verbose = True)

Consider the following list of lists:

l = [[11, 12, 13], [21, 22, 23, 24], [31, 32]]
l_flat = [11, 12, 13, 21, 22, 23, 24, 31, 32]
lengths = [3, 4, 2]

Counting from 0, the 4th part of l_flat is 1st part of the 1st list. Hence, index_of_which_bin(lengths, 4) will output 1, 1.