What's a good way to access a 2D list with a "flattened" index?

Viewed 1163

Suppose we have a 2D, 3x3 list:

lst =  [[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]]

What would be the recommended way to access list elements with a "flattened" index?

I mean, I have only a number from 0 to 9. How do I access the corresponding list[i][j] element?

5 Answers

you can use divide and modulo to compute 2D coordinates (only works if all sub lists have the same size)

def get_item(lst,i):    
    return(lst[i//len(lst[0])][i%len(lst[0])])

note that flattening the list to access items is O(n**2) complexity, whereas this method is O(1).

If you have variable sublist length, you have to find the sublist first, and you cannot use a division. But once the proper sublist is found, you have O(1) access:

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

def get_item(lst,i):
    current = 0
    for sublist in lst:
        index = i - current
        if 0 <= index < len(sublist):
            return sublist[index]
        current += len(sublist)
    raise Exception("index out of range {}".format(i))

print(get_item(lst,0),get_item(lst,6),get_item(lst,9))

prints: 1 7 10

Here is a short command however, less efficient:

lst = [[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]]

#Get the index no. 5 from 2d list
sum(lst, [])[5]

Output:

6

You can use itertools.chain:

from itertools import chain

lst = [[1, 2, 3],
         [4, 5, 6],
         [7, 8, 9]]

for i, x in enumerate(chain(*lst)):
    print(i, x)

0 1
1 2
2 3
3 4
4 5
5 6
6 7
7 8
8 9

So to leverage your flattened index:

# you can just match the index explicitly
for i, x in enumerate(chain(*lst)):
    if i==5:
        break

x
# 6

# or you can store in a list 
flattened = list(chain(*lst))

flattened[5]
# 6

You can use numpy.reshape to generate a 1D array from a 2D array.

import numpy

in_list = [[1,2,3],[4,5,6],[7,8,9]]
out_list = numpy.reshape(in_list, 9).tolist()

print(in_list)
#[[1,2,3],[4,5,6],[7,8,9]]

print(out_list)
#[1,2,3,4,5,6,7,8,9]

Then you could just access your resized array as a 1D array

print(out_list[2])
#2

For a rectangular matrix, you can use divmod to ger the row and column indexes Access Complexity: O(1):

matrix = [[1, 2, 3,],
          [10,11,12],
          [21,22,23]]

r,c = divmod(6,len(matrix[0]))
value = matrix[r][c]
print(value) # 21

For an irregular list of lists, as long as none of the sizes change, you can build an indirection list and use it to convert a flat position into row and column indexes Access Complexity: O(1) Indexing complexity O(N) (N=total number of elements):

matrix = [[1, 2, 3, 4 ,5],
          [10,11,12],
          [21,22,23,24]]

indx = [ (r,c) for r,row in enumerate(matrix) for c in range(len(row)) ]
pos   = 8
r,c   = indx[pos]
value = matrix[r][c]

print(value) # 21

Once the indx list is created, accessing elements will be practically as fast as the rectangular matrix method (integer division and modulo).

Note, you would have to recreate the indirection list (indx) any time you change the size of the list or if any of its rows.

If your matrix is changed (in size) more often than you access it using flat positions, you could create a function to return the indexes Access Complexity: O(R):

def rowCol(m,p):
    for r,row in enumerate(m):
        if p<len(row): return r,p
        p-=len(row)

r,c = rowCol(matrix,8)
value = matrix[r][c]
print(value) 

This will be a lot slower to access than the prepared indx variable so you may want to use hybrid approach where you clear the index after each change and only rebuild it on demand (when the function is called)

Another alternative that will build the index faster but will be a bit slower to access would be to create an index with only the cumulative sizes of the matrix's rows. You can then use a binary search algorithm to find the row and a simple subtraction to find the column Access Complexity: O(log(R)) Indexing complexity: O(R)

matrix = [[1, 2, 3,],
          [10,11,12,13,14],
          [21,22,23,24]]

from itertools import accumulate
from bisect import bisect_left
indx = [0]+list(accumulate(map(len,matrix)))

pos   = 8
r     = bisect_left(indx,pos+1)-1
c     = pos - indx[r]
value = matrix[r][c]
print(pos,r,c,value) # 21

This is slower to access than the full coordinate indexing but it will consume less memory and will still be much faster than the iterative function (which essentially does the same thing sequentially)

Related