Komo
Komo

Reputation: 13

Index of an element in a nested list

I'm struggling with an exercise for a few days. Given was following nested list:

[1, [5, 62, 6], 4, [99, [100, 200, 600, [1000, [2000]]]], [74, 41, 16], 7, [8], [[[400]]]]

And this function body:

def find_element(liste, find, index = 0):

I have to find an element in the nested list and the function should return the exact index of the found element, for example [1,0] for 5 or [3, 1, 3, 1, 0] for 2000. The function has to be recursive.

My problem is the function has to return false if the element is not in the list.

This is my code:

def find_element(liste, find, index = 0):
    indexList = []

    if len(liste) == index:
        return indexList

    if liste[index] == find:
        indexList.append(index)
    else:
        if type(liste[index]) == list:
            indexList.extend(find_element(liste[index], find))
        if indexList:
            indexList.insert(0, index)
        else:
            indexList.extend(find_element(liste, find, index + 1))

    return indexList

I tried a second function that returns false when the list is empty or an if condition if the index is 0 and the indexList is empty, but all I got were a RecursionError or a TypeError.

Upvotes: 1

Views: 1024

Answers (3)

Mulan
Mulan

Reputation: 135367

I agree generators are a nice fit for this problem. I would separate the program logic into two separate functions, dfs and find_element -

def dfs(ls, r = []):
  if isinstance(ls, list):
    for (i, v) in enumerate(ls):
      yield from dfs(v, [*r, i])
  else:
    yield (r, ls)

def find_element(ls, q):
  for (k, v) in dfs(ls):
    if v == q:
      return k
  return None
print(find_element(input, 5))
# [1, 0]

print(find_element(input, 2000))
# [3, 1, 3, 1, 0]

print(find_element(input, 999))
# None

Or you could fix your original program using a fourth parameter, r = [] -

def find_element(ls, q, i = 0, r = []):
  if i >= len(ls):
    return None
  elif isinstance(ls[i], list):
    return find_element(ls[i], q, 0, [*r, i]) \
      or find_element(ls, q, i + 1, r)
  elif ls[i] == q:
    return [*r, i]
  else:
    return find_element(ls, q, i + 1, r)
print(find_element(input, 5))
# [1, 0]

print(find_element(input, 2000))
# [3, 1, 3, 1, 0]

print(find_element(input, 999))
# None

Upvotes: 1

tclarke13
tclarke13

Reputation: 395

Ajax1234's answer works, but if you need something a little more simple this may be better:

def find_idx(input_list, elem):
  for i in range(len(input_list)):
    if isinstance(input_list[i], list):
      result = find_idx(input_list[i], elem)
      if result:
        return [i] + result
    elif input_list[i] == elem:
        return [i]

  return False

input_list = [1, [5, 62, 6], 4, [99, [100, 200, 600, [1000, [2000]]]], [74, 41, 16], 7, [8], [[[400]]]]

print(find_idx(input_list, 2000))
# Output: [3, 1, 3, 1, 0]  

 

This is basically a DFS (https://en.wikipedia.org/wiki/Depth-first_search). If you think of your data structure as a tree, your list entries are nodes since they themselves can contain other lists, just as a node in a tree can point to other nodes. The magic is in returning False if nothing was found at the very end of the method, but recursively searching all sublists before you get to that point. Also, you have to check whether your list entry is itself a list, but this is just an analogy to the fact that a tree can have nodes that do point to other nodes, and nodes that do not (leaf nodes, or plain old numbers in your case).

Upvotes: 2

Ajax1234
Ajax1234

Reputation: 71461

You can use recursion with a generator:

def find_element(l, elem):
  def get_elem(d, c = []):
    for i, a in enumerate(d):
       if a == elem:
          yield c+[i]
       elif isinstance(a, list):
          yield from get_elem(a, c+[i])
  return False if not (r:=list(get_elem(l))) else r[0]

data = [1, [5, 62, 6], 4, [99, [100, 200, 600, [1000, [2000]]]], [74, 41, 16], 7, [8], [[[400]]]]
print(find_element(data, 2000))

Output:

[3, 1, 3, 1, 0]

Upvotes: 1

Related Questions