Can this list comprehension be improved?

Question

Let’s say I have this list:

input_list = [{'a': 1, 'b': 2}, {'a': 3, 'b': 5}]

I want to create output_list based on input_list:

output_list = []
for dic in input_list:
    new_dic = {}
    ab_sum = sum([dic['a'], dic['b']])
    if ab_sum % 2 == 0:
        new_dic['c'] = ab_sum
        new_dic['d'] = ab_sum ** 2
        new_dic['e'] = ab_sum ** 4
        output_list.append(new_dic)

Result:

[{'c': 8, 'd': 64, 'e': 4096}]

The actual dictionary is way bigger and this gets messy. The more readable solution would be to use list comprehension:

output_list = [{'c': ab_sum,
                'd': (ab_sq:= ab_sum **2),
                'e': ab_sq **2}
                for dic in input_list 
                if (ab_sum:=sum([dic['a'], dic['b']])) % 2 == 0]

This seems inconsistent as I assign to variables both in the filter and within the dictionary.
I would like to know if there is a more elegant solution to these types of problems, or I am overthinking it?

Asked By: sierra_papa

||

Source

Answer 1

i think what you’ve done is as far as you can go with a list comprehension, maybe youll find its more readable if you use some functions, and you can choose whether you want to use filters/maps or comprehensions

def sum_ab(dic):
  return dic['a'] + dic['b']

def is_ab_sum_even(dic):
  return sum_ab(dic) % 2 == 0

def get_cde(dic):
  ab_sum = sum_ab(dic)
  return {"c": ab_sum, "d": ab_sum **2, "e": ab_sum **4}

out_list = [get_cde(dic) for dic in in_list if is_ab_sum_even(dic)]

OR, you can do maps/filters

out_list = list(map(get_cde, filter(is_ab_sum_even, in_list)))

On the assigning to a variable matter, i think it might be slightly on the overthinking end because doing a + b, even a million times for a computer is super easy, but so is storing a single variable and continuously overwriting it

Answered By: rymanso

Answer 2

Here is an alternative version, using a helper function, and taking advantage of the walrus operator :=, introduced in Python 3.8 (I changed the variable names just to make the code easier to read):

def get_cde_from(pair):
    ab = sum(pair.values())
    if not (ab % 2):
        return {*zip('cde', [ab, ab * ab, ab ** 4])}


pairs = [
    {'a': 1, 'b': 2},
    {'a': 3, 'b': 5}
]

triplets = []

for pair in pairs:
    if cde := get_cde_from(pair):
        triplets.append(cde)

print(triplets)

The output of the above code is the following:

[{'c': 8, 'd': 64, 'e': 4096}]

You can also populate triplets (output_list, in your sample code) with a comprehension combined with the walrus operator:

triplets = [cde for pair in pairs if (cde := get_cde_from(pair))]

Another option is to use a generator instead:

def gen_cde_from(pairs):
    for pair in pairs:
        ab = sum(pair.values())
        if not (ab % 2):
            yield dict(zip('cde', [ab, ab * ab, ab**4]))


pairs = [
    {'a': 1, 'b': 2},
    {'a': 3, 'b': 5}
]

triplets = [*gen_cde_from(pairs)]

print(triplets)

Which will also result in:

[{'c': 8, 'd': 64, 'e': 4096}]

Answered By: accdias

Can this list comprehension be improved?

Question:

Answers: