Ad

Convert Flattened 1 Dimensional Indices To 2 Dimensional Indices

Say I have a list of lists, e.g:

x = [[0,1,2,3],[4,5],[6,7,8,9,10]]

And I have the 'flat' indices of the elements I wish to target, i.e, the indices of the elements I want to select from the list if it were flattened into a 1d list:

flattened_indices = [0,1,4,9]

                  # #     #         #
flattened_list = [0,1,2,3,4,5,6,7,8,9,10]

How do I convert the 1.d. indices into 2.d. indices that would allow me to recover the elements from the original nested list? I.e. in this example:

2d_indices = [(0,0), (0,1), (1,0), (2,3)]
Ad

Answer

Here is a way to do that:

from bisect import bisect
import itertools

# Accumulated sum of list lengths
def len_cumsum(x):
    return list(itertools.accumulate(map(len, x)))

# Find 2D index from accumulated list of lengths
def find_2d_idx(c, idx):
    i1 = bisect(c, idx)
    i2 = (idx - c[i1 - 1]) if i1 > 0 else idx
    return (i1, i2)
# Test
x = [[0, 1, 2, 3], [4, 5], [6, 7, 8, 9, 10]]
indices = [0, 4, 9]
flattened_list = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
c = len_cumsum(x)
idx_2d = [find_2d_idx(c, i) for i in indices]

print(idx_2d)
>>> [(0, 0), (1, 0), (2, 3)]

print([x[i1][i2] for i1, i2 in idx_2d])
>>> [0, 4, 9]

If you have many "flat" indices, this is more effective than iterating the nested list for each index.

Ad
source: stackoverflow.com
Ad