Question

I have a function that has a dictionary as input and a value n.
Each item in the dictionary is a set with one or more values.
The function should sort the dictionary keys and it should extract and return “n”values.
This function will be executed very often therefore I am trying to optimize it. Any suggestions?

def select_items(temp_dict, n):
  """Select n items from the dictionary"""
  res = []
  sort_keys = sorted(temp_dict.keys())
  count = 0

  for key in sort_keys:
    for pair in temp_dict[key]:
      if count < n:
        res.append(pair)
        count += 1
      else:
        return res

  return res

In this code I have a count and “if statement” to control the number of selected values. Is there a way to optimize this code by using some function in itertools or something else?

Answer 1

The answers given so far don’t follow the user’s spec.

The data is a dictionary of sequences, and the desired result is a list of the first n elements of the dictionary values taken in sorted order by key.

So if the data is:

{1: [1, 2, 3], 2: [4, 5, 6]}

then, if n = 5, the result should be:

[1, 2, 3, 4, 5]

Given that, here’s a script which compares the original function with a (slightly) optimised new version:

from timeit import timeit

def select_items_old(temp_dict, n):
  res = []
  sort_keys = sorted(temp_dict.keys())
  count = 0
  for key in sort_keys:
    for pair in temp_dict[key]:
      if count < n:
        res.append(pair)
        count += 1
      else:
        return res
  return res

def select_items_new(data, limit):
    count = 0
    result = []
    extend = result.extend
    for key in sorted(data.keys()):
        value = data[key]
        extend(value)
        count += len(value)
        if count >= limit:
            break
    return result[:limit]

data = {x:range(10) for x in range(1000)}

def compare(*args):
    number = 1000
    for func in args:
        name = func.__name__
        print ('test: %s(data, 12): %r' % (name, func(data, 12)))
        code = '%s(data, %d)' % (name, 300)
        duration = timeit(
            code, 'from __main__ import %s, data' % name, number=number)
        print ('time: %s: %.2f usec/pass\n' % (code, 1000000 * duration/number))

compare(select_items_old, select_items_new)

Output:

test: select_items_old(data, 12): [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1]
time: select_items_old(data, 300): 163.81 usec/pass

test: select_items_new(data, 12): [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1]
time: select_items_new(data, 300): 67.74 usec/pass