Question

I have data like

[2, 2, 2, 2, 2, 3, 13, 113]

which I then want to sort into separate lists by keys generated by myself. In fact I want to generate all possible lists.

Some examples:

values: [2, 2, 2, 2, 2, 3, 13, 113]
keys:   [0, 0, 1, 2, 1, 3,  3,   1]
sublists: [2, 2], [2, 2, 113], [2], [3, 13]

values: [2, 2, 2, 2, 2, 3, 13, 113]
keys:   [0, 1, 0, 0, 0, 1,  1,   0]
sublists: [2, 2, 2, 2, 113], [2, 3, 13]

values: [2, 2, 2, 2, 2, 3, 13, 113]
keys:   [2, 3, 0, 0, 4, 4,  1,   3]
sublists: [2, 2], [13], [2], [2, 113], [2, 3]

All possible keys are generated by

def generate_keys(prime_factors):
    key_size = len(prime_factors) - 1
    key_values = [str(i) for i in range(key_size)]

    return list(itertools.combinations_with_replacement(key_values, \
        len(prime_factors)))

Then I thought I could use the keys to shift the values into the sublists. That’s the part I’m stuck on. I thought itertools.groupby would be my solution but upon further investigation I see no way to use my custom lists as keys for groupby.

How do I split my big list into smaller sublists using these keys? There may even be a way to do this without using keys. Either way, I don’t know how to do it and looking at other Stack Overflow questions has eben in the ballpark but not exactly this question.

Answer 1

This does what you want:

def sift(keys, values):
    answer = collections.defaultdict(list)
    kvs = zip(keys, values)
    for k,v in kvs:
        answer[k].append(v)
    return [answer[k] for k in sorted(answer)]

In [205]: keys = [0, 0, 1, 2, 1, 3,  3,   1]

In [206]: values = [2, 2, 2, 2, 2, 3, 13, 113]

In [207]: sift(keys,values)
Out[207]: [[2, 2], [2, 2, 113], [2], [3, 13]]

Explanation:

collections.defaultdict is a handy dict-like class that lets you define what should happen in the event that a key doesn’t exist in the dictionary that you’re trying to manipulate. For example, in my code, I have answer[k].append(v). We know that append is a list function, so we know that answer[k] should be a list. However, if I was using a conventional dict and I tried to append to the value of a non-existent key, I would have gotten a KeyError as follows:

In [212]: d = {}

In [213]: d[1] = []

In [214]: d
Out[214]: {1: []}

In [215]: d[1].append('one')

In [216]: d[1]
Out[216]: ['one']

In [217]: d
Out[217]: {1: ['one']}

In [218]: d[2].append('two')
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
/Users/USER/<ipython-input-218-cc58f739eefa> in <module>()
----> 1 d[2].append('two')

KeyError: 2

This was only made possible because I defined answer = collections.defaultdict(list). If I had defined answer = collections.defaultdict(int), I would gotten a different error – one that would tell me that int objects don’t have an append method.

zip on the other hand takes two lists (well actually, it takes at least two iterables), lets call them list1 and list2 and returns a list of tuples in which the ith tuple contains two objects. The first is list1[i] and the second is list2[i]. If list1 and list2 are of unequal length, len(zip(list1, list2)) would be the smaller value among len(list1) and len(list2) (i.e. min(len(list1), len(list2)).

Once I’ve zipped keys and values, I want to create a dict such that maps a value from keys to a list of values from values. This is why I used a defaultdict, so that I wouldn’t have to check for the existence of a key in it before I appended to its value. If I had used a conventional dict, I would have had to do this:

answer = {}
kvs = zip(keys, values)
for k,v, in kvs:
    if k in answer:
        answer[k].append(v)
    else:
        answer[k] = [v]

Now that you have a dict (or a dict-like object) that maps values from keys to lists of ints that share the same key, all you need to do is get the lists which are the values of answer in sorted order, sorted by the keys of answer. sorted(answer) gives me a list of all of answers keys in sorted order.

Once I have this list of sorted keys, all I have to do is get their values, which are lists of ints, and put all those lists into one big list and return that big list.

… annnnnd Done! Hope that helps