CODEWITHSUNDEEP
pythonpython-3.x
SteveMcManaman
SteveMcManaman

Reputation: 433

How can I group objects of a class, based on consecutive attributes?

We have a class Shift:

class Shift:
    def __init__(self, station, week, dayofweek, date):
        self.station = station
        self.week = week
        self.dayofweek = dayofweek
        self.date = date

and we have a sorted list (by 'date' of these shifts) of instances from this class. I want to group those shifts which have consecutive dates together and store the groups in a list at the end.

Input: [Shift(date=1),Shift(date=2),Shift(date=4),Shift(date=5),Shift(date=6)]

Output:[[Shift(date=1),Shift(date=2)],[Shift(date=4),Shift(date=5),Shift(date=6)]]

Upvotes: 0

Views: 42

Answers (1)

Mark
Mark

Reputation: 92440

You need to keep track of the last Shift seen and then decide whether to add the next one to the last group or start a new group. Here's a paired down version:

class Shift:
    def __init__(self, date):
        self.date = date

    # some helpers to make other code clearer
    # and direct comparisons possible
    def __repr__(self):
        return f"Shift({self.date})"

    def __lt__(self, other):
        return self.date < other.date

    def __sub__(self, other):
        return self.date - other.date

l = [Shift(1),Shift(6),Shift(5),Shift(4),Shift(2),Shift(9),Shift(7), Shift(0)] 

def groupNear(l):
    if not l:
        return l
    group = []
    current = l[0]

    for s in sorted(l):
        if s - current <= 1:
            group.append(s)
        else:
            yield group
            group = [s]
        current = s 
    yield group

list(groupNear(l))

Output:

[[Shift(0), Shift(1), Shift(2)],
 [Shift(4), Shift(5), Shift(6), Shift(7)],
 [Shift(9)]]

Upvotes: 1

Related Questions