CODEWITHSUNDEEP
python
Brian
Brian

Reputation: 151

How to compress this if else statement

if word not in word_map.keys():
    set1 = set()
    set1.add(line_no)                
    word_map[word] = set1
else:
    set1 = set()
    set1.add(line_no)
    set1.update(word_map[word])
    word_map[word] = set1

Is there anyway I can compress this into not needing an if? I have to check if the key exists so I know if I need to add what is already in set key or not, if I just have

set1.update(word_map[word])

And the key does not exist, it will error out.

Upvotes: 0

Views: 111

Answers (2)

poke
poke

Reputation: 387825

First of all, you can compress it by removing duplicate code like this:

set1 = set()
set1.add(line_no)
if word in word_map.keys():
    set1.update(word_map[word])
word_map[word] = set1

Then you should think about what set.update does. It basically adds all the elements from one set into the current one. So you can reorder your code like this too:

if word in word_map.keys():
    set1 = word_map[word]
else:
    set1 = set()
set1.add(line_no)
word_map[word] = set1

We can then make this simpler by using dict.get to remove the if completely:

set1 = word_map.get(word, set())
set1.add(line_no)
word_map[word] = set1

And then, we can also use dict.setdefault instead to make sure that we have a set, and just add to that directly:

set1 = word_map.setdefault(word, set())
set1.add(line_no)

As we now use the set that’s inside the dictionary, we don’t need to write back its value to the dictionary (we have a reference). And we could even make it a single line then:

word_map.setdefault(word, set()).add(line_no)

Upvotes: 4

David Heffernan
David Heffernan

Reputation: 613053

Perhaps you are looking for this:

set1 = set()
set1.add(line_no)                
if word in word_map.keys():
    set1.update(word_map[word])
word_map[word] = set1

Although, on the assumption that word_map is a dict, the if should be written like this:

if word in word_map:

Upvotes: 1

Related Questions