Can this list comprehension be improved?-CodePudding

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?

CodePudding user response：

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

CodePudding user response：

Here is an alternative version, using a helper function, and taking advantage of the walrus operator :=, introduced in Python 3.8:

def get_cde(d):
    ab = sum(d.values())
    if not (ab % 2):
        return dict(zip('cde', [ab, ab * ab, ab ** 4]))
    return None


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

output_list = []

for d in input_list:
    if cde := get_cde(d):
        output_list.append(cde)

print(output_list)

The output of the above code is the following:

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

Another option is to use a generator instead:

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


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

output_list = list(gen_cde(input_list))

print(output_list)

Which will also result in:

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