Python: linear __getitem__ for a pair of list of lists

I currently have a class which stores a list of lists. The inner lists are not of the same length. I made the class subscriptable with the following code (possibly not the best way of doing this, and perhaps overly fancy).

class MyClass:
    def __init__(self):
        #
        self.instructions = []

        # for demo purposes
        self.instructions.append([0, 1, 2])
        self.instructions.append([3, 4, 5, 6])
        self.instructions.append([7, 8])

    def __getitem__(self, ind):
        if ind >= 0:
            iterator = self.instructions.__iter__()
            compare = int.__gt__
            inc = int.__add__
        else:
            iterator = reversed(self.instructions)
            compare = int.__le__
            inc = int.__sub__

        s = 0
        for tp in iterator:
            L = len(tp)
            if compare(inc(s, L), ind):
                return tp[ind-s]
            else:
                s = inc(s, L)
        else:
            raise IndexError('index out of range')

This works. For instance

>>> x = MyClass()
>>> x[5]
5
>>> x[-5]
4

Now, I need to modify the class so it now stores two list of lists. The two lists are instructions and annotations, and both have the same length. But len(instructions[i]) does not have to equal len(annotations[i]).

class NewClass:
    def __init__(self):
        #
        self.instructions = []
        self.annotations = []

        # for demo purposes
        self.instructions.append([0, 1, 2])
        self.instructions.append([5, 6, 7, 8])
        self.instructions.append([12, 13])
        
        self.annotations.append([3, 4])
        self.annotations.append([9, 10, 11])
        self.annotations.append([14, 15, 16])

    def __getitem__(self, ind):
        pass

I want to make this subscriptable, with the order of elements oscillating between the instructions sublists and the annotations sublists. The demo data indicates the subscripting order. I want

>>> y = NewClass()
>>> y[9]
9
>>> y[-4]
13

What's an efficient way of doing this?

I could write a solution where I alternatively iterate through the two sublists. But I feel like I am straying far from the correct solution. I am also looking for a non-for-loop solution as I am dealing with long lists.

The standard balance between storage cost and runtime cost for random access into the (non-stored) concatenation of several arrays is to store a table of the offsets of the beginning of each list (i.e., the sum of the length of every list before it) and use binary search on that table:

import itertools
import bisect
class Index:
def __init__(self,ll):
self.ll = ll
self.oo = list(itertools.accumulate(map(len,ll), initial=0))
def __getitem__(self, i):
if i < 0:
i += self.oo[-1]
j = bisect.bisect(self.oo,i)
if not 0 < j <= len(self.ll):
raise IndexError
return self.ll[j-1][i-self.oo[j-1]]
def __iter__(self):
return itertools.chain.from_iterable(self.ll)
# Example:
i = Index(
[[9,1,7],
[3,0],
[],
[4,4,4,2]]
)
assert i[4]==0 and i[8]==2

The j-1 is because the initial 0 causes an i of 0 to be assigned the insertion point of 1. You can omit the ,initial=0 (and in fact the last element of self.oo as well) at the cost of more complicated code for the edge/error cases. __iter__ is provided because it is asymptotically faster than indexing with successive integers, each of which must be subjected to the binary search even though usually the same sublist will be found.

Obviously extending this to support the interleaving of two lists (of equal length) is trivial: sum the interleaved lengths and then use divmod(j-1,2) to obtain the index into a list and the selection between the two lists (respectively).

While your implementation is nice, I would like to share my own way for iterating using chain.from_iterable. Because basically we're chaining the items whether from the beginning or at the end.

For one list:

The only part that needs explanation is map(reversed, reversed(self.instructions)). We not only need to reverse the whole list, but also the individual sublists.

from itertools import chain

class MyClass:
    def __init__(self):
        self.instructions = [
            [0, 1, 2],
            [3, 4, 5, 6],
            [7, 8],
        ]

    def __getitem__(self, ind):
        if ind >= 0:
            chunks = self.instructions
            range_parameter = ind + 1
        else:
            chunks = map(reversed, reversed(self.instructions))
            range_parameter = abs(ind)

        iterator = chain.from_iterable(chunks)

        try:
            for _ in range(range_parameter):
                n = next(iterator)
        except StopIteration:
            raise IndexError("index out of range")

        return n

x = MyClass()
print(x[5])
print(x[-5])

For two lists:

Since you said we need to oscillate, zip is the right tool for that. When ind is positive it's straightforward. We zip them and use chain.from_iterable two times because otherwise it gives us individual sub-lists.

If ind is negative, we need two reversed() before zipping. One for outer lists, and one for sublists.

from itertools import chain

class MyClass:
    def __init__(self):
        self.instructions = [
            [0, 1, 2],
            [5, 6, 7, 8],
            [12, 13],
        ]

        self.annotations = [
            [3, 4],
            [9, 10, 11],
            [14, 15, 16],
        ]

    def __getitem__(self, ind):
        if ind >= 0:
            chunks = zip(self.instructions, self.annotations)
            range_parameter = ind + 1
        else:
            chunks = zip(
                map(reversed, reversed(self.annotations)),
                map(reversed, reversed(self.instructions)),
            )
            range_parameter = abs(ind)

        iterator = chain.from_iterable(chain.from_iterable(chunks))

        try:
            for _ in range(range_parameter):
                n = next(iterator)
        except StopIteration:
            raise IndexError("index out of range")

        return n

x = MyClass()
print(x[9])
print(x[-4])

here is my idea, for the 2 list version

first we define a funtion to alternate between lists

from itertools import zip_longest
def alternate(*iterables: "list[list[Any]]") -> "Iterator[list[Any]]":
for group in zip_longest(*iterables):
for it in group:
if it is not None:
yield it

like that in can alternate between any number of your lists

now for the class

class MyClass:
def __init__(self):
self.instructions = [
[0, 1, 2],
[5, 6, 7, 8],
[12, 13],
]
self.annotations = [
[3, 4],
[9, 10, 11],
[14, 15, 16],
]
def __len__(self):
return sum(map(len,alternate(self.instructions, self.annotations)))
def __getitem__(self, index:int ):
size = len(self)
if index >= 0:
if index >= size:
raise IndexError(index)
else:
new_index = size + index
if 0 <= new_index < size:
index = new_index
else:
raise IndexError(index)
for item in alternate(self.instructions, self.annotations):
n = len(item)
if 0 <= index < n:
return item[index]
index -= n
raise RuntimeError("this should not happens")
test = MyClass()
print(test[9])
print(test[-4])

I defined a len which is the sum of all the length of the sublist, and use that in the __getitem__ to first determinate if the input index is valid or not, and for the negative case first calculate the positive version of the index and check if its correct.

With that done, and given that the sublist aren't of the same size, then simple going in a loop for each of the sublist and if the index fall in its range return that otherwise subtract the size of that sublist and go to the next.

and because python is cool like that you get for free that your class is iterable (just by having and __getitem__ that take interger values) and reversible (by also having a __len__)

>>> list(test)
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
>>> list(reversed(test))
[16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0]

alternatively you can use the nth recipe

from itertools import islice
def nth(iterable, n, default=None):
"Returns the nth item or a default value"
return next(islice(iterable, n, None), default)   

(or from the more_itertools third party library)

and then alongside itertools.chain you can change the final loop in getitem for return nth(chain.from_iterable(alternate(self.instructions, self.annotations)), index) given that we already check that the index fall in range

Here is my approach:

import itertools


class NewClass:
    def __init__(self):
        #
        self.instructions = []
        self.annotations = []

        # for demo purposes
        self.instructions.append([0, 1, 2])
        self.instructions.append([5, 6, 7, 8])
        self.instructions.append([12, 13])
        
        self.annotations.append([3, 4])
        self.annotations.append([9, 10, 11])
        self.annotations.append([14, 15, 16])
        
    def __iter__(self):
        zipped = itertools.zip_longest(self.instructions, self.annotations, fillvalue=[])
        for sub_lists in zipped:
            yield from itertools.chain.from_iterable(sub_lists)

    def __getitem__(self, key):
        flat = list(self)
        return flat[key]

Notes

• I created the __iter__ method, which let the caller iterates through the object like this:

    x = NewClass()
  for e in x:
  print(e)
  

• The __getitem__ is built upon __iter__

• About the zipped data: Conceptually, you can view this as

    [
  [[0, 1, 2], [3, 4]],          # This is a sub_lists
  [[5, 6, 7, 8], [9, 10, 11]],  # a sub_lists
  ...
  ]
  

• The expression itertools.chain.from_iterable(sub_lists) basically flatten a sub_lists from [[0, 1, 2], [3, 4]] to [0, 1, 2, 3, 4]

• This solution works for an arbitrary number of lists, not just 2.

Update

I fixed up the __getitem__ to handle negative index at the cost of performance. I'm too lazy to create a solution which is more efficient.

If the two concatenated lists are the same size, you can use something like this:

div, mod = divmod(ind, 2)
if mod:
return get_item(second_list, div)
else:
return get_item(first_list, div)