How to sort list of student records by categories in Python

Viewed 722

Suppose that I have nested list that looks like this

# Student records
data_list=[ # person, subject, grade
   ["John", "Physics", 5], ["John", "PC", 7], ["John", "Math", 8], 
   ["Mary", "Physics", 6], ["Mary", "PC", 10], ["Mary", "Algebra", 7], 
   ["Helen", "Physics", 7], ["Helen","PC", 6], ["Helen", "Algebra", 8], 
   ["Helen", "Analysis", 10], ["Bill", "PC", 10], ["Bill", "Analysis", 6], 
   ["Bill", "Math", 8], ["Bill", "Biology", 6], ["Michael", "Analysis", 10]
]

How can I create a code that prints what subject did each student attend to? The code output should look like this:

# Subjects taken by person
[["John", "Physics", "PC", "Math"],
["Mary", "Physics", "PC", "Algebra"],
...]

I would prefer a solution that uses only Python lists.

2 Answers

How about using a dictionary to collect subjects per person:

from collections import defaultdict

subjects = defaultdict(list)
for record in data_list: 
    subjects[record[0]].append(record[1])
# subjects
defaultdict(list,
            {'John': ['Physics', 'PC', 'Math'],
             'Mary': ['Physics', 'PC', 'Algebra'],
             'Helen': ['Physics', 'PC', 'Algebra', 'Analysis'],
             'Bill': ['PC', 'Analysis', 'Math', 'Biology'],
             'Michael': ['Analysis']})

This can then be converted to a list of lists:

subjects_list = [[x] + y for x, y in subjects.items()]
# subjects_list
[['John', 'Physics', 'PC', 'Math'],
 ['Mary', 'Physics', 'PC', 'Algebra'],
 ['Helen', 'Physics', 'PC', 'Algebra', 'Analysis'],
 ['Bill', 'PC', 'Analysis', 'Math', 'Biology'],
 ['Michael', 'Analysis']]

EDIT As requested by the OP, here is a solution using only lists:

subject_list = []  # output
persons = []       # track unique persons
for record in data_list:
    if record[0] not in persons:
        persons.append(record[0])           # Track new person
        subject_list.append([record[0]])    # Add new person to output
        subject_list[-1].append(record[1])  # Add subject for person
    else:
        subject_list[persons.index(record[0])].append(record[1])
# subject_list
[['John', 'Physics', 'PC', 'Math'],
 ['Mary', 'Physics', 'PC', 'Algebra'],
 ['Helen', 'Physics', 'PC', 'Algebra', 'Analysis'],
 ['Bill', 'PC', 'Analysis', 'Math', 'Biology'],
 ['Michael', 'Analysis']]

EDIT 2 You can generalise this approach to sort your data_list by any index (e.g. person, ...) filtering out a category (e.g. subject, grade, ...)

def groupby(data, index, category):
    """Sort list of records by index and category
    """
    output = []
    indices = []
    for record in data:
        if record[index] not in indices:
            indices.append(record[index])
            output.append([record[index]])
            output[-1].append(record[category])
        else:
            output[indices.index(record[index])].append(record[category])
    return output

This will allow you to list subjects per person like this:

# index 0 -> person
# category 1 -> subject
subject_list = groupby(data_list, 0, 1)

Or you could list grades per person like this:

# index 0 -> person
# category 2 -> grade
grade_list = groupby(data_list, 0, 2)
# grad_list
[['John', 5, 7, 8],
 ['Mary', 6, 10, 7],
 ['Helen', 7, 6, 8, 10],
 ['Bill', 10, 6, 8, 6],
 ['Michael', 10]]

You can then get the number of subjects taken per person or the average grade like this:

import statistics

subjects_taken = [len(x) - 1 for x in subject_list]
average_grade = [statistics.mean(x[1:]) for x in grade_list]

Putting everything together gives you:

persons = [x[0] for x in subject_list]
final_list = list(zip(persons, subjects_taken, average_grade))

# final_list
[('John', 3, 6.666666666666667),
 ('Mary', 3, 7.666666666666667),
 ('Helen', 4, 7.75),
 ('Bill', 4, 7.5),
 ('Michael', 1, 10)]

Here is a function you could try:

def consolidate(data):
    # list of names from  data_list
    names = list(set([i[0] for i in data])) 
    consolidated = [] # empty list to get the sublists

    # Looping for each name
    for name in names:
        subList = [name] # sublist for each name
        for i in data: # loop for each item within data_list
            if i[0] == name: # if name is the same, append the subject
                subList.append(i[1])
        consolidated.append(subList) # append sublists in consolidated list
    return consolidated

Call function:

consolidate(data_list)

Output:

[['Mary', 'Physics', 'PC', 'Algebra'], ['Helen', 'Physics', 'PC', 'Algebra', 'Analysis'],
['John', 'Physics', 'PC', 'Math'], ['Bill', 'PC', 'Analysis', 'Math', 'Biology'],
['Michael', 'Analysis']]
Related