How to Create a Custom Nested Index for Multidimensional Lists?

by

Problem Formulation

Have you ever wanted to get the index of an item in a nightmarish nested list of lists of…, with different lengths and datatypes?

# Multi-dimensional list:
headache_list = [0, [[1, "", 2, "three"], 
                     [4, 5.0, [6, "seven"]], 
                     [8, 9], 
                     []]]

# What is the index of element 8?
print(headache_list.index(8))

Output:

ValueError: 8 is not in list

Well, you have found to the right place!

💬 Follow me through this step-by-step tutorial, and you’ll end up with an elegant functional solution to create a multi-dimensional indexing solution for nested Python lists. You don’t need any external modules.

Naive Solution for 2D List

Things would be easier if we had to deal with nested lists of the same known length and depth.

If we have a simple two-dimensional “matrix” stored in the list, we can hard-code indexes or use list comprehensions and generator expressions in five steps:

  • Step 1. Determine all rows with query element
  • Step 2. Take first row with query element
  • Step 3. Determine row index
  • Step 4. Determine column index
  • Step 5. Print result

Here’s the Python code that outlines the exact implementation of each step:

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

###########
# What is the index of query in my_list?
###########

# Step 1. Determine all rows with query element
rows = [row for row in my_list if query in row]

# Step 2. Take first row with query element
r = rows[0]

# Step 3. Determine row index
i = my_list.index(r)

# Step 4. Determine column index
j = r.index(query)

# Step 5. Print result
print(f'List: {my_list}')
print(f'Index of element {query} in list is ({i}, {j})')
print(f'my_list[{i}][{j}] =', my_list[i][j])

But we are finxters, and we are not satisfied with easy but incomplete solutions.

So let’s think about how to approach the problem of having lists of different depths, lengths, and data types… Recursion!

Recursive Solution for General Problem

Solution Overview:

We will recursively compare the object of which we want to find the index with each element of the list of lists until we have a match.

We will use enumerate() to get the index of the iterable element we will search on.

  • If there is a match, we return an empty list in which we will insert the indexes that led us there.
  • If there is no match, we return a None object.

We will know that we have reached the end of each branch of the list if the next element on which we want to iterate is not iterable.

We will use try-except to catch the TypeError when using enumerate() with a non-iterable argument.

Problem: What if we have an empty or single-character string?

It is iterable, and we would enter an infinite loop when iterating over it since an enumerate object is not None:

>>> type(enumerate("")) 
<class 'enumerate'>

To solve this, we will use a condition to check if the object we would iterate over next is a string and if its length is <= 1. If it evaluates as True, we will return None, and we’ll go to the next potential branch.

If we finish every possible branch without a match, we will gracefully unwind and return None.

If we had a match, we would have returned an empty list [] which is not None, so the condition to insert every index recursively in the first position of the list would kick in, and we would return a list of indices to show off our awesome skills.

Here is my solution with some examples:

def nested_index(item, chaos):

    # Return empty list to fill with indexes if item found
    if item == chaos:
        # Found item!
        # Return an empty list that is a 'not None' 
        # object which will end recursion.
        return []

    # If 'chaos' is an empty or single char string, with no 
    # match, we're out of luck! We don't want to fall into a 
    # pit of endless recursion, so we return None.
    if isinstance(chaos, str) and len(chaos) <= 1:
        return None

    # If 'chaos' is not iterable, this is a dead-end!
    try:
        # Here's the beauty of recursion!
        for index, heap in enumerate(chaos):
            index_list = nested_index(item, heap)
            # Time to check our heap...
            if index_list is not None:
                # We've found it! Time to unwind the results :D
                index_list.insert(0, index)
                return index_list
    except TypeError:
        pass

    # We haven't found what we were looking for, too bad...
    return None

Let’s apply this nested_index() function to a number of examples to understand how it is used:

headache_list = [0, [[1, "", 2, "three"], 
                     [4, 5.0, [6, "seven"]], 
                     [8, 9], 
                     []]]

print(nested_index(8, headache_list))
print(nested_index(5.0, headache_list))
print(nested_index([8, 9], headache_list))
print(nested_index("seven", headache_list))
print(nested_index("v", headache_list))
print(nested_index("", headache_list))
print(nested_index([], headache_list))
print(nested_index(headache_list, headache_list))
print(nested_index("finxter", headache_list))

The output is:

[1, 2, 0]
[1, 1, 1]
[1, 2]
[1, 1, 2, 1]
[1, 1, 2, 1, 2]
[1, 0, 1]
[1, 3]
[]
None

Here you can see the function working step by step:

This solution is based on Daniel Kullmann’s answer in this SO post.

To keep improving your Python skills, we recommend you check out our free email academy — we have cheat sheets too!