Преобразование одномерных индексов в двухмерные

Question

Скажем, у меня есть список списков, например:

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

И у меня есть «плоские» индексы элементов, на которые я хочу ориентироваться:

indices = [0,1,4,9]

, т. Е.индексы элементов, которые я хочу выбрать из списка, если он был сведен в одномерный список:

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

Как мне преобразовать 1.d.индексы в 2.d.индексы, которые позволили бы мне восстановить элементы из исходного вложенного списка?Т.е. в этом примере:

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

Answer 1

Вот способ сделать это:

from bisect import bisect

# Accumulated sum of list lengths
def len_cumsum(x):
    c = [len(lst) for lst in x]
    for i in range(1, len(c)):
        c[i] += c[i - 1]
    return c
# As pointed out by @tobias_k, this can actually be just done as:
# 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]

Если у вас много «плоских» индексов, это более эффективно, чем повторение вложенного списка для каждого индекса.

Answer 2

Полагаю, вы могли бы поместить эти индексные пары в dict, а затем просто сослаться на dict из indices в конце и создать новый список:

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

indices = [0,4,9]

idx_map = {x: (i, j) for i, l in enumerate(x) for j, x in enumerate(l)}

result = [idx_map[x] for x in indices]
print(result)

В результате:

[(0, 0), (1, 0), (2, 3)]

Но это не оптимально, поскольку его квадратичное время выполнения создает idx_map.Решение @ jdehesa , использующее bisect, намного более оптимально.