Python: Iterator, Generator

 

Sequence Protocol

To make a custom data type have Sequence Protocol, you need to implement "__iter__" method. Even though "__getitem__" method is enough for now, you should implement "__iter__" method as well for later compatibility.

 

A classic iterator

Below is a classic implementation - not a rolemodel - of an iterator.

import re
import reprlib


class Sentence:

    def __init__(self, text):
        self.text = text
        self.words = re.compile('\w+').findall(text)

    def __repr__(self):
        return f"Sentence{reprlib.repr(self.text)}"

    def __iter__(self):
        return SentenceIterator(self.words)


class SentenceIterator:

    def __init__(self, words):
        self.words = words
        self.index = 0

    def __next__(self):
        try:
            word = self.words[self.index]
        except IndexError:
            raise StopIteration()
        self.index += 1
        return word

    def __iter__(self):
        return self


s = Sentence('Hello my name is Python')
print(repr(s))
for x in s:
    print(x)

SentenceIterator only has two methods - "__next__" and "__iter__'. We have a seperate class only for iterator to ensure that iter(sentence obj) returns a new iterator.

 

A more pythonic way would be using generator.

class Sentence:

    def __init__(self, text):
        self.text = text
        self.words = re.compile('\w+').findall(text)

    def __repr__(self):
        return f"Sentence{reprlib.repr(self.text)}"

    def __iter__(self):
        yield from self.words

However, above is still inefficient in terms of memory - let's improve it with 'finditer()'.

class Sentence:

    def __init__(self, text):
        self.text = text

    def __repr__(self):
        return f"Sentence{reprlib.repr(self.text)}"

    def __iter__(self):
        for match in re.compile('\w+').finditer(self.text):
            yield match.group()

Or, we can replace __iter__ method with a 'generator expression'.

    def __iter__(self):
        return (match.group() for match in re.compile('\w+').finditer(self.text))

 

 

Example: An arithmetcial progression 

def aritprog_gen(begin, step, end=None):
    result = type(begin+step)(begin)
    endless = True if end is None else False
    index = 0
    while endless or result < end:
        yield result
        index += 1
        result = begin + step * index

 

 

Generator functions in STL

(1) filter(predicate, it)

(2) filterfalse(preciate, it)

(3) isslice(it, start, stop, step=1)

import itertools

test = 'aAbBcCdDeE'

for x in filter(str.islower, test):
    print(x, end=' ')
print()

for x in itertools.filterfalse(str.islower, test):
    print(x, end=' ')
print()

for x in itertools.islice(test, 2):
    print(x, end=' ')
print()

 

(4) accumulate(it, [func])

import itertools

nums = [5, 7, 0, 7, 1, 1, 9]

for x in itertools.accumulate(nums, max):
    print(x, end=' ')
print()

(5) enumerate(it, start=0)

(6) map(func, it1, it2, ...)

(7) strmap(func, it)

import itertools
import operator

test = 'aAbBcCdDeEfF'

for x in map(operator.mul, range(11), range(11)):
    print(x, end=' ')
print()

for x in map(str.lower, test):
    print(x, end=' ')
print()

for x in itertools.starmap(operator.mul, enumerate(range(11), 1)):
    print(x, end=' ')
print()

 

(8) chain(it1, it2, ...)

(9) chain.from_iterable(it)

(10) product(it1, ..., itN, repeat=1)

(11) zip(it1, it2, ...)

(12) zip_longest(it1, it2, ..., fillvalue=None)

for x in itertools.chain('ABC', range(1, 3)):
    print(x, end=' ')
print()


for x in itertools.product('ABC', range(1, 3)):
    print(x, end=' ')
print()

 

(13) combinations(it, out_len)

(14) combinations_with_replacement(it, out_len)

(15) count(start=0, step=1)

(16) permutation(it, out_len=None)

(17) repeat(item, [times])

 

(18) groupby(it, key=None)

(19) reversed(seq)

(20) tee(it, n=2)

LeetCode: 2092. Find All People With Secret Solution

from itertools import groupby, tee

class Solution:
    def findAllPeople(self, n: int, meetings, firstPerson: int):
        parents = [i for i in range(n)]
        parents[firstPerson] = 0
        
        def find(a):
            if a == parents[a]:
                return a
            parents[a] = find(parents[a])
            return parents[a]
        
        def union(a, b):
            pa, pb = find(a), find(b)
            if pa == 0 or pb == 0:
                parents[pa] = parents[pb] = 0
            else:
                parents[pb] = pa
        
        for _, grp in groupby(sorted(meetings, key=lambda x: x[2]), key=lambda x: x[2]):
            grp1, grp2 = tee(grp, 2)
            for x, y, t in grp1:
                union(x, y)
            for x, y, t in grp2:
                px = find(x)
                if px != 0:
                    parents[x] = x
                    parents[y] = y
        return [i for i in range(n) if find(i) == 0]