Question

For a homework assignment, we were required to create a function that took in two inputs (list and a string) and that returned a list. I would like to write this function using list comprehension, but have been running into an annoying issue. Here is he original function I would like to re-write:

index = [['keyword1', ['url1', 'url2', 'url3']], ['keyword2', ['url4', 'url5']]...]

def lookup(index, keyword):
    result = []
    for entry in index:
        if entry[0] == keyword:
            result += entry[1]
    return result

These were my attempts at using list comprehension:

def lookup(index, keyword):
    return [entry[1] for entry in index if entry[0] == keyword]

def lookup(index, keyword):
    return [ulist for key, ulist in index if i == keyword]

and finally…

def lookup(index, keyword):
    return [j for j in (entry[1] for entry in index if entry[0] == keyword)]

The problem here is that it returns the desired urls in a list inside of a list, like so:

[['url1', 'url2', 'url3']]

instead of the desired format of:

['url1', 'url2', 'url3']

I thought of just doing something like this (adding [0] to the end of the return statement so it gives me the inner list):

def lookup(index, keyword):
        return [j for j in (entry[1] for entry in index if entry[0] == keyword)][0]

But it just doesn’t seem like the right way. Am I trying to use list comprehension unnecessarily (my understanding was that this situation was well suited for list comprehension)? Any suggestions?

Answer 1

This would be the most straightforward way:

def lookup(index,keyword):
    return [x for entry in index if entry[0]==keyword for x in entry[1]]

The nested list comprehension syntax isn’t what you might initially expect. The key is to realize that the use of the two for loops is not two separate list comprehensions, but a single list comprehension with nested loops. That is actually why it works. It is like you were writing

for entry in index:
  if entry[0]==keyword:
    for x in entry[1]:
      yield x

The order in list comprehensions is the same as in regular loops, except for the yielded value, which comes first.