Joe Iddon
Joe Iddon

Reputation: 20414

Indexing a 2D List with another List

If I have a list such as:

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

and another list which contains a [row, column] index for an element in the list, so for example:

index = [2, 1]

Obviously I can get the element at index with:

lst[index[0]][index[1]]

However, I find that this syntax is quite clunky and doesn't sit well in a big block of code. It would also become worse with a higher dimensional list.

My question: is there an easier way to do this index?

It would seem that something like:

lst[index]

would be more readable but this isn't how Python works so isn't an option.

Upvotes: 3

Views: 2371

Answers (3)

niemmi
niemmi

Reputation: 17263

You could just create a simple function that iterates over the index. For every element in index just fetch item from object and assign that as a new object. Once you have iterated over the whole index return current object. As @EricDuminil noted it works with dicts and all other objects that support __getitem__:

def index(obj, idx):
    for i in idx:
        obj = obj[i]

    return obj

LST = [[1,2,3], [4,[5],6], [{'foo': {'bar': 'foobar'}},8,9]]
INDEXES = [[2, 2], [1, 1, 0], [2, 0, 'foo', 'bar']]

for i in INDEXES:
    print('{0} -> {1}'.format(i, index(LST, i)))

Output:

[2, 2] -> 9
[1, 1, 0] -> 5
[2, 0, 'foo', 'bar'] -> foobar

Upvotes: 2

Eric Duminil
Eric Duminil

Reputation: 54223

Since you're working with a 2D-list, it might be a good idea to use numpy. You'll then simply need to define index as a tuple. Index 3 would be out of range, though:

>>> import numpy as np
>>> a = np.array([[1,2,3], [4,5,6], [7,8,9]])
>>> index = (1, 2)
>>> a[index]
6

The method you're looking for is called Array#dig in Ruby:

[[1,2,3], [4,5,6], [7,8,9]].dig(1, 2)
# 6

but I couldn't find any plain Python equivalent.

Upvotes: 2

ividito
ividito

Reputation: 336

The method you used is probably the simplest way, but for better readability you could go for something like

i = index[0]
j = index[1]
x = lst[i][j]

One way or another, you need to split your index list into the 2 values. If you want to declutter your main block of code, you can write a function to handle this, but that would hardly be "easier".

EDIT: As suggested below, tuple unpacking is an even better option

i, j = index
x = lst[i][j]

Upvotes: 1

Related Questions