For flat indexing into a nested structure without flattening that structure.

multi_level_index.py

import bisect
import numpy as np
import collections.abc

class TwoLevelIndex(collections.abc.Sequence):
    def __init__(self, coll):
        counts = [len(thing) for thing in coll]
        self.cum_counts = np.cumsum(counts)
    def __getitem__(self, idx):
        seg_idx = bisect.bisect(self.cum_counts, idx)
        offset = idx - (self.cum_counts[seg_idx-1] if seg_idx > 0 else 0)
        return seg_idx, offset
    def __len__(self):
        return self.cum_counts[-1]
    def __iter__(self):
        for i in range(len(self)):
            yield self[i]

class MultiLevelIndex(collections.abc.Sequence):
    def __init__(self, coll, levels):
        assert levels >= 2
        self.deeper = []
        counts = []
        for thing in coll:
            if levels > 3:
                deeper_index = MultiLevelIndex(thing, levels-1)
            else:
                deeper_index = TwoLevelIndex(thing)
            self.deeper.append(deeper_index)
            counts.append(len(deeper_index))
        self.cum_counts = np.cumsum(counts)
    def __getitem__(self, idx):
        seg_idx = bisect.bisect(self.cum_counts, idx)
        offset = idx - (self.cum_counts[seg_idx-1] if seg_idx > 0 else 0)
        return (seg_idx, *self.deeper[seg_idx][offset])
    def __len__(self):
        return self.cum_counts[-1]
    def __iter__(self):
        for i in range(len(self)):
            yield self[i]

test_multi_level_index.py

import random

import multi_level_index

coll = [[[0,1,2,3], [4,5,6], [7,8,9]],[[10, 11], [12, 13], [14, 15], [16]],[[17, 18, 19], [20]]]
m_index = multi_level_index.MultiLevelIndex(coll, 3)
for x,y,z in random.sample(m_index, len(m_index)):
    print(x, y, z, coll[x][y][z])

x,y,z = m_index[20]
print(x,y,z)

It's a recursive implementation, so, you know, it's probably not very fast with deeply nested structures.