Grouping a dictionary of lists based on <size of list> % key and the first value in the list

Question

How to best optimize this code ?

Problem Description:

Input(counts) is an array of \$n\$ integers, where each counts[i], (\$0\le i< n\$) denotes the total number of elements in a particular group that element i belongs to. For example, if counts = [3, 3, 3, 3, 3, 1, 3], then there are three groups; elements 0, 1, 2, 3, 4, and 6 are in one of two 3-element groups, and element 5 is in a 1-element group.

A group is valid if all the elements in the group have minimal ID numbers. In other words, a group of size \$k\$ must contain the \$k\$ smallest ID numbers belonging to a group of that size with respect to the smallest element ID in the group. For example, if counts = [3, 3, 3, 3, 3, 1, 3], then the grouping [0, 1, 2], [3, 4, 6], and [5] is valid; however, the grouping [0, 1, 4], [2, 3, 6], and [5] is not valid because the group [0, 1, 4] does not contain the three smallest element IDs for the set of element IDs belonging to 3-element groups (i.e., \${0,1,2,3,4,6}\$).

It is guaranteed that a valid grouping always exists for the given input array.

VALID GROUPING:

0 1 2
3 4 6
5

INVALID GROUPING:

0 1 4
2 3 6
5

Sample Input (1):

counts = [2,1,1,2,1]

Sample Output (1):

0 3 
1 
2 
4

Sample Input (2):

counts = [4,2,4,5,5,4,4,5,5,2,5,7,1,7,1,7,1,7,1,7,1,7,1,7,5,5,5,5,5]

Sample Output (2):

0 2 5 6
1 9
3 4 7 8 10 
11 13 15 17 19 21 23
12
14
16
18
20
22
24 25 26 27 28

Code:

from collections import OrderedDict

def groupCount(counts):
    group_dict = OrderedDict()
    for i in range(len(counts)):
        if counts[i] in group_dict.keys():
            group_dict[counts[i]].append(i)
        else:
            group_dict[counts[i]] = [i]
    op_list = []
    for key in group_dict:
        prnt_count = 0
        temp_list = []
        while len(group_dict[key])!=0:
            prnt_count += 1
            temp_list.append(group_dict[key].pop(0))
            if prnt_count % key == 0:
                op_list.append(temp_list)
                temp_list = []
    op_list.sort(key=lambda x: x[0])
    for value in op_list:
        print(*value, sep=' ')

Answer 1

1. Review

1. There's no docstring. What does groupCount do? The text in the post explains it very clearly, so using this would be a good start.

2. groupCount has two responsibilities: (i) it collates the input into groups, and (ii) it prints the collated results. The trouble with combining responsibilities like this is that it makes it hard to reuse the code. If you need to use the groups in other code, how do you get them? If you want to write automatic test cases, how do you do that? It would be better to split the grouping and printing code into separate functions. (This is the "single responsibility principle".)

3. This code iterates over the indexes in the sequence counts and then looks up the item using counts[i].

   for i in range(len(counts)):
       # code using i and counts[i]
   

   Python provides the function enumerate for simultaneously iterating over items and their indexes:

   for i, count in enumerate(counts):
       # code using i and count
   

4. For collating items based on their keys, it is handy to use collections.defaultdict, like this:

   groups = defaultdict(list)
   for i, count in enumerate(counts):
       groups[count].append(i)
   

5. To split a sequence into groups of length \$n\$, there's a well-known trick that's described in the documentation for zip, where it says:

   This makes possible an idiom for clustering a data series into n-length groups using zip(*[iter(s)]*n). This repeats the same iterator n times so that each output tuple has the result of n calls to the iterator. This has the effect of dividing the input into n-length chunks.

6. There is no need to pass a key function to sort. Sequences are compared lexicographically by their items, and because we know that all the items are different, the comparison will never have to look past the first item.

2. Revised code

from collections import defaultdict

def group_by_count(counts):
    """Given a sequence of group sizes for each item, return a list of
    groups, each group being a tuple of indexes of items.

    >>> group_by_count([3, 3, 3, 3, 3, 1, 3])
    [(0, 1, 2), (3, 4, 6), (5,)]

    """
    groups = defaultdict(list)
    for i, count in enumerate(counts):
        groups[count].append(i)
    result = []
    for count, group in groups.items():
        result.extend(zip(*[iter(group)] * count))
    return sorted(result)

(Note the example in the docstring: this can be automatically checked using the doctest module.)