Python: Removing elements from list that is a value in a dictionary?

Question

I have the following dictionary:

a = dict()

a["A"] = [1,2,3,4,5]

a["B"] = [1,2]

I want to iterate through each value (in this case in the form of a list) in the dictionary and if the number of elements in the list is greater than 2, I want to remove all the elements after the second element.

For example, after I alter the values, my dictionary should be:

a["A"] = [1,2] # [3,4,5] are removed from this value

a["B"] = [1,2]

I tried:

del a["A"][3:len(a[2])]

But this returned:

[3, 4, [1, 2, 3]]

Two questions: 1) Why is the output of my attempt the way it is? 2) How would I correctly implement what I am trying to do?

Thanks!

Answer 1

The simplest way of implementing what you want is "writing it literally":

for value in dictionary.values():
    del value[2:]

Which could be read as:

for every value in the dictionary, remove all elements after the second.

Note that:

• This solution does not create a new dictionary, it modifies the original dictionary in-place
• This solution does not create new lists. It modifies the current lists in-place.

In certain circumstances you may want to create a new dictionary and/or avoid mutating the values of the dictionary, in this case you should avoid using del and instead re-assign the sliced values:

dictionary[key] = value[:2]

---

It seems also the most efficient:

In [1]: %%timeit
   ...: a = {'A': [1,2,3,4,5], 'B': [1,2]}
   ...: for value in a.values():
   ...:     del value[2:]
   ...: 
1000000 loops, best of 3: 663 ns per loop
In [2]: %%timeit
   ...: a = {'A': [1,2,3,4,5], 'B': [1,2]}
   ...: for key, value in a.items():
   ...:     a[key] = value[:2]
   ...: 
1000000 loops, best of 3: 920 ns per loop
In [3]: %%timeit
   ...: a = {'A': [1,2,3,4,5], 'B': [1,2]}
   ...: for value in a.values():
   ...:     if len(value) > 2:
   ...:         del value[2:]
   ...: 
1000000 loops, best of 3: 796 ns per loop

In [4]: %%timeit
   ...: a = {'A': [1,2,3,4,5], 'B': [1,2]}
   ...: a = {k:v[:2] for k,v in a.items()}
   ...: 
1000000 loops, best of 3: 1.15 us per loop

If most of the values are of length 2 or shorter, then adding an if len(value) > 2: to slice only when necessary might give a slight increase in performances. If you expect most values to be longer than 2 then it only increases overhead.

---

Regarding your first question: Doing del a["A"][3:len(a[2])] does not give [3, 4, [1, 2, 3]] as output. You get a KeyError: 2 . You have done something else if you obtained that output.

Answer 2

Use this function:

def trim_dictionary(d, trim=2):
    return {k:v[:trim] for k,v in d.iteritems()}

Example:

>>> trim_dictionary({'A':[1,2,3,4], 'B':[5,6], 'C':[7,8,9], 'D':[10]})
{'A': [1, 2], 'B': [5, 6], 'C': [7, 8], 'D': [10]}

Answer 3

Use a dict comprehension:

>>> {k: [item for item in v if item <= 2] for k, v in a.iteritems()}
{'A': [1, 2], 'B': [1, 2]}

To modify the original dict:

for k, v in a.iteritems():
...     v[:] = [item for item in v if item <= 2]
...     
>>> a
{'A': [1, 2], 'B': [1, 2]}

To keep just the first two items in each list:

for k in a:
    a[k] = a[k][:2]  #or del a[k][2:]
...     
>>> a
{'A': [1, 2], 'B': [1, 2]}

Answer 4

Instead of

del a["A"][3:len(a[2])]

Use

a["A"] = a["A"][:2]

or

del a["A"][2:]

Answer 5

Just get the (max) first two elements.

a["A"] = [1,2,3,4,5]
a["B"] = [1,2]
a["C"] = [1]
for key in a:
    a[key] = a[key][:2]

If you have a lot of data it might be faster to check the length of the list.

for key, value in a.items():
    if len(value) > 2:
        a[key] = value[:2]

Method 3 is using del :

for value in a.values():
    del value[2:]

All examples tested on Python 3.