flat_map_python.py

from functools import reduce
import timeit

class Test:
    x = 2

def id(x):
    return x

def check_time(f, arg):
    n = 50
    return 1000 * (timeit.timeit(lambda: loop(f, id, arg), number=n) / n)  # ms

def measure(f):
    xs = [list(range(10)) for _ in range(100)]
    print("a)", f"{check_time(f, xs):.4f}", "ms")

    xs = [list(range(10)) for _ in range(10000)]
    print("b)", f"{check_time(f, xs):.4f}", "ms")

    xs = [[Test() for _ in range(10)] for _ in range(10000)]
    print("c)", f"{check_time(f, xs):.4f}", "ms")


def loop(f, g, arg):
    ys = f(g, arg)
    for y in ys:
        pass

def flat_map(f, xs):
    ys = []
    for x in xs:
        ys.extend(f(x))
    return ys

print("double for loop")
measure(flat_map)

flat_map = lambda f, xs: [y for ys in xs for y in f(ys)]
print("comprehension")
measure(flat_map)

flat_map = lambda f, xs: (f(y) for ys in xs for y in ys)
print("generator")
measure(flat_map)

flat_map = lambda f, xs: reduce(lambda a, b: a + b, map(f, xs))
print("map reduce")
measure(flat_map)

flat_map = lambda f, xs: [f(x) for x in reduce(lambda a, b: a + b, xs)]
print("comprehension reduce")
measure(flat_map)

flat_map = lambda f, xs: sum(map(f, xs), [])
print("map sum")
measure(flat_map)

flat_map = lambda f, xs: [f(x) for x in sum(xs, [])]
print("comprehension sum")
measure(flat_map)