Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

How can I use if/else in a dictionary comprehension?

People also ask

How do you write if-else in comprehension?

Worth mentioning that you don't need to have an if-else condition for both the key and the value. For example, {(a if condition else b): value for key, value in dict. items()} will work. @JeremyWeirich You don't need to have an if-else for either of them if you don't want to.

Can you put an if statement inside a dictionary Python?

Python dictionary comprehension is a method for transforming one dictionary into a new dictionary. Using an if-else in Python dict comprehension is one of the ways to create a new list. It works as a condition to filter out dictionary elements for a new dictionary.

How do you check if a key is in a dictionary?

How do you check if a key exists or not in a dictionary? You can check if a key exists or not in a dictionary using if-in statement/in operator, get(), keys(), handling 'KeyError' exception, and in versions older than Python 3, using has_key(). 2.

Can you use list comprehension with dictionaries?

Dictionary data structure lets you query the data using key very efficiently. The idea of comprehension is not just unique to lists in Python. Dictionaries, one of the commonly used data structures in data science, can also do comprehension.


You've already got it: A if test else B is a valid Python expression. The only problem with your dict comprehension as shown is that the place for an expression in a dict comprehension must have two expressions, separated by a colon:

{ (some_key if condition else default_key):(something_if_true if condition
          else something_if_false) for key, value in dict_.items() }

The final if clause acts as a filter, which is different from having the conditional expression.


@Marcin's answer covers it all, but just in case someone wants to see an actual example, I add two below:

Let's say you have the following dictionary of sets

d = {'key1': {'a', 'b', 'c'}, 'key2': {'foo', 'bar'}, 'key3': {'so', 'sad'}}

and you want to create a new dictionary whose keys indicate whether the string 'a' is contained in the values or not, you can use

dout = {"a_in_values_of_{}".format(k) if 'a' in v else "a_not_in_values_of_{}".format(k): v for k, v in d.items()}

which yields

{'a_in_values_of_key1': {'a', 'b', 'c'},
 'a_not_in_values_of_key2': {'bar', 'foo'},
 'a_not_in_values_of_key3': {'sad', 'so'}}

Now let's suppose you have two dictionaries like this

d1 = {'bad_key1': {'a', 'b', 'c'}, 'bad_key2': {'foo', 'bar'}, 'bad_key3': {'so', 'sad'}}
d2 = {'good_key1': {'foo', 'bar', 'xyz'}, 'good_key2': {'a', 'b', 'c'}}

and you want to replace the keys in d1 by the keys of d2 if there respective values are identical, you could do

# here we assume that the values in d2 are unique
# Python 2
dout2 = {d2.keys()[d2.values().index(v1)] if v1 in d2.values() else k1: v1 for k1, v1 in d1.items()}

# Python 3
dout2 = {list(d2.keys())[list(d2.values()).index(v1)] if v1 in d2.values() else k1: v1 for k1, v1 in d1.items()}

which gives

{'bad_key2': {'bar', 'foo'},
 'bad_key3': {'sad', 'so'},
 'good_key2': {'a', 'b', 'c'}}

In case you have different conditions to evaluate for keys and values, @Marcin's answer is the way to go.

If you have the same condition for keys and values, you're better off with building (key, value)-tuples in a generator-expression feeding into dict():

dict((modify_k(k), modify_v(v)) if condition else (k, v) for k, v in dct.items())

It's easier to read and the condition is only evaluated once per key, value.

Example with borrowing @Cleb's dictionary of sets:

d = {'key1': {'a', 'b', 'c'}, 'key2': {'foo', 'bar'}, 'key3': {'so', 'sad'}}

Assume you want to suffix only keys with a in its value and you want the value replaced with the length of the set in such a case. Otherwise, the key-value pair should stay unchanged.

dict((f"{k}_a", len(v)) if "a" in v else (k, v) for k, v in d.items())
# {'key1_a': 3, 'key2': {'bar', 'foo'}, 'key3': {'sad', 'so'}}

It also might be worth mentioning that If only statements put the if at the end:

{_ for _ in iterable if True}