does python have a partition_by function?

Question

is there some established recipe for partitioning a list according to some predicate? What I want to do is like an itertools groupby, but the predicate is going to be some arbitrarily complicated function (and not just a key). For example, imagine a list of students, and each student has a list of courses, and I want to group by those with courses in common. So it would be something like:

def coursework_intersection(a,b):
     return set(a['courses']).intersection(b['courses'])

list_of_lists = partition_by(coursework_intersection, students) 

Answer 1

If you want what Joran says in comments then it's inherently an Omega(n^2) worst case running time because that's the size of the output in the worst case (where coursework_intersection always returns true). So let's just bite the bullet:

def associated_with(func, seq):
    for item in seq:
        yield item, (other for other in seq if func(item, other))

Note that the input is a sequence, not an iterable, since this is a multi-pass algorithm.

That could be optimized to call func half as many times, if we're allowed to assume that it's a symmetric function, although the cost is more memory use. It could also be optimized a bit into a one-liner return ( (item, (other for other in seq if func(item, other))) for item in seq) , but I judge that not the most readable way to introduce the code ;-)

Answer 2

from collections import defaultdict

def group_by_classes(students):
    result = defaultdict(list)
    for student in students:
        result[set(student["courses"])].append(student)
    return result

which will result in a list of students for each unique set of classes (ie every occupied vertex of the class hypercube).