List as Value of a Dictionary with Multiple Values in Each List

Question

I am trying to solve the following problem using Python:

Given the following student list:

students = ['Tom', 'Kitty', 'Jessie', 'Chester', 'Curie', 'Darwing', 'Nancy', 'Sue', 'Peter', 'Andrew', 'Karren', 'Charles', 'Nikhil', 'Justin', 'Astha']

I need to assign the students into 3 groups. The group identity of each student is specified by the integer at the same index position in a list named assignment:

assignment = [2, 3, 3, 1, 3, 2, 3, 1, 1, 1, 2, 1, 3, 2, 2]

Using comprehension, I need to generate dictionary with Group number as the key and list of students belonging to that group as the value, ie. the expected output is like this:

{'Group 1': ['Chester', 'Sue', 'Peter', 'Andrew', 'Charles'],
 'Group 2': ['Tom', 'Darwing', 'Karren', 'Justin', 'Astha'],
 'Group 3': ['Kitty', 'Jessie', 'Curie', 'Nancy', 'Nikhil']}

My first step is to zip the 2 lists together to generate a list of tuples:

zip(students,assignment)

I generated an empty dictionary to store the group membership: cdict = {}

My logic is: Check the integer (group number) in each tuple. If the number does not exist in the dictionary keys, create a new one and add the name corresponding to that group number as well. If it exists, add the value into the value list in the dictionary.

I came up with the following code:

groupdict = {'Group{}'.format(group):[name] for name,group in zip(students,assignment) if group not in groupdict.keys()}

However, it gave me the following result:

{'Group 2': ['Astha'], 'Group 3': ['Nikhil'], 'Group 1': ['Charles']}

I tried to include the else statement:

groupdict = {'Group{}'.format(group):[name] for name,group in zip(students,assignment) if group not in groupdict.keys() else groupdict[group].append(name)} 

but it gives me a SyntaxError.

How should I amend my comprehension code in order to store a list of values rather than only the latest value? Is there a way to complete this other than using setdefault() function or importing libraries? I am expected to only include for, if, else and zip() in my code. Thanks for help in advance! :)

Answer 1

Using dict.setdefault

Ex:

students = ['Tom', 'Kitty', 'Jessie', 'Chester', 'Curie', 'Darwing', 'Nancy', 'Sue', 'Peter', 'Andrew', 'Karren', 'Charles', 'Nikhil', 'Justin', 'Astha']
assignment = [2, 3, 3, 1, 3, 2, 3, 1, 1, 1, 2, 1, 3, 2, 2]
result = {}
for k,v in zip(students,assignment):
    result.setdefault(f"Group {v}", []).append(k)

Output:

{'Group 1': ['Chester', 'Sue', 'Peter', 'Andrew', 'Charles'],
 'Group 2': ['Tom', 'Darwing', 'Karren', 'Justin', 'Astha'],
 'Group 3': ['Kitty', 'Jessie', 'Curie', 'Nancy', 'Nikhil']}

Answer 2

Another solution, if you want to do it with dict-comprehension (involves using itertools.groupby and sorting (Rakesh's answer is more effective in this regard)):

from itertools import groupby


students = ['Tom', 'Kitty', 'Jessie', 'Chester', 'Curie', 'Darwing', 'Nancy', 'Sue', 'Peter', 'Andrew', 'Karren', 'Charles', 'Nikhil', 'Justin', 'Astha']
assignment = [2, 3, 3, 1, 3, 2, 3, 1, 1, 1, 2, 1, 3, 2, 2]

out = {'Group {}'.format(group): [s for s, _ in g] for group, g in groupby(sorted(zip(students, assignment), key=lambda k: k[1]), lambda k: k[1])}
print(out)

Prints:

{'Group 1': ['Chester', 'Sue', 'Peter', 'Andrew', 'Charles'], 
 'Group 2': ['Tom', 'Darwing', 'Karren', 'Justin', 'Astha'], 
 'Group 3': ['Kitty', 'Jessie', 'Curie', 'Nancy', 'Nikhil']}

Answer 3

using defaultdict

from collections import defaultdict
result = defaultdict(list)
students = [
    'Tom', 'Kitty', 'Jessie', 'Chester', 'Curie', 'Darwing', 'Nancy', 'Sue',
    'Peter', 'Andrew', 'Karren', 'Charles', 'Nikhil', 'Justin', 'Astha'
]
assignment = [2, 3, 3, 1, 3, 2, 3, 1, 1, 1, 2, 1, 3, 2, 2]
for k, v in zip(students, assignment):
    result[f"Group {v}"].append(k)
print(result)

output

defaultdict(<class 'list'>, {'Group 2': ['Tom', 'Darwing', 'Karren', 'Justin', 'Astha'], 'Group 3': ['Kitty', 'Jessie', 'Curie', 'Nancy', 'Nikhil'], 'Group 1': ['Chester', 'Sue', 'Peter', 'Andrew', 'Charles']})