Nested list comprehension with if statement

Question

I'm trying to get my head around nested list comprehension and have read the excellent explanation here.

The problem I'm having translating is that I've an if clause in my inner loop and I can't see how to apply this to the func() step as I'm loosing the counter I get from enumerate() when I go from nested loops to list comprehension.

nested_list = [[{'a': 1, 'b': 2}, {'c': 3, 'd': 4}], [{'a': 5, 'b': 6}, {'c': 7, 'd': 8}]]
new_list = []
for c, x in enumerate(nested_list):
    for d, y in enumerate(x):
        if d == 1:
            new_list.append(y)
print(new_list)
[{'c': 3, 'd': 4}, {'c': 7, 'd': 8}]

Nested list comprehension might look something like

new_list = [if ??? y
               for x in nested_list
                   for y in x]

...but I can't see/think how to get the clause as I have no counter under the nested list comprehension.

Is there a way of achieving this or should I stick to the nested loops approach?

List comprehensions use the *same nesting order*, left to right. Move the `if` to the same location, don't put it in the value expression. — Martijn Pieters, Sep 27 '18 at 12:27
And you dropped the `enumerate()` function call, now you don't *have* a `d` variable to test with. — Martijn Pieters, Sep 27 '18 at 12:27
@MartijnPieters yes the loss of `d` in the second enumerate is the crux of my problem. I'll give yours and t@Jean-FrancoisFabre suggestion a whirl, thanks. — slackline, Sep 27 '18 at 12:31
@Sayse I believe list comprehension is a more "pythonic" way of doing things and will give cleaner more maintainable code. — slackline, Sep 27 '18 at 12:32
@slackline: while a list comprehension is indeed a better option for a loop that builds a list by simply appending, your loop is doing a lot more work than is strictly necessary. — Martijn Pieters, Sep 27 '18 at 12:37

Jean-François Fabre · Accepted Answer · 2018-09-27T12:41:23.110

12

you can rewrite this as follows:

new_list = [y for x in nested_list for d, y in enumerate(x) if d == 1]

loops are in the "natural" order, and the condition is in the end. Include y only if d==1

one possible result (because dicts aren't ordered):

[{'c': 3, 'd': 4}, {'c': 7, 'd': 8}]

Note that in your case, it's simpler and more efficient (O(n) vs O(n**2)) to write:

new_list = [x[1] for x in nested_list]

the only difference is that if x is too short the latter code will break so maybe test length:

new_list = [x[1] for x in nested_list if len(x)>1]

edited Sep 27 '18 at 12:41

answered Sep 27 '18 at 12:29

Jean-François Fabre

137,073
23
153
219

Thanks that's really useful, especially the shorter `x[1]` way of doing it as in my actual use case the nested lists will always be the same length. – slackline Sep 27 '18 at 12:38

Martijn Pieters · Answer 2 · 2018-09-27T12:39:18.413

as I'm loosing the counter I get from enumerate() when I go from nested loops to list comprehension.

Then why did you drop it? A list comprehension can still use all the same iterators as a for loop can.

I've an if clause in my inner loop and I can't see how to apply this

Just put it in the same place. List comprehensions follow the same nested order for both for loops and if filters; your loop can be mapped directly by removing the : colons, and moving the part in the .append(...) call to the front:

new_list = [y   # from new_list.append(y)
    for c, x in enumerate(nested_list)   # removed the :
        for d, y in enumerate(x)         # removed the :
            if d == 1]                   # removed the :

However, you can drop the first enumerate(), because you don't use c. You can also drop the second loop and enumerate(), because you could use indexing instead:

new_list = [x[1] for x in nested_list]

You don't need to loop over all values in x here, that's not a very efficient method to get to that one value.

If your issue is that the x list might be shorter and not have an element at index 1, filter those nested lists on length first:

new_list = [x[1] for x in nested_list if len(x) > 1]

Thanks, having only been using Python for a few months I'm still learning and hadn't realised that you can leave the `if` statements as is. Very clear explanation, cheers — slackline, Sep 27 '18 at 13:11

score 1 · Answer 3 · answered Sep 27 '18 at 12:48

1

You want the second element for each array in the list, so you can do something like this:

new_list = [nested_list[i][1] for i in range(len(nested_list))]

better still

new_list  = [x[1] for x in nested_list]

answered Sep 27 '18 at 12:48

Bruno Vermeulen

2,970
2
15
29

score 0 · Answer 4 · answered Sep 27 '18 at 13:08

0

print ([ y for c, x in enumerate(nested_list) for d, y in enumerate(x) if d == 1])

answered Sep 27 '18 at 13:08

Mahesh

1,583
13
18

2

While this might answer the authors question, it lacks some explaining words and links to documentation. Raw code snippets are not very helpful without some phrases around it. You may also find [how to write a good answer](https://stackoverflow.com/help/how-to-answer) very helpful. Please edit your answer. – hellow Sep 27 '18 at 13:47

Nested list comprehension with if statement

4 Answers4

Related