Sum and count a stream of numbers coming from a generator expression

Viewed 145

I can sum a stream of numbers coming from a generator expression

education = '0 1 0 0 0'.split()
salary = [int(s) for s in '50 120 0 40 60'.split()] # 0 for missing data
# compute the sum of known low education salaries, aka les
total_les = sum(s for e, s in zip(education, salary) if e=='0' and s>0)

now I'd like to compute the mean of the known salaries, in my example it's 150/3

I could get a list of low education salaries

 list_of_les = [s for e, s in zip(education, salary) if e=='0' and s>0]
 mean_les = sum(list_of_les)/len(list_of_les)

or I could count again my stream

 count = sum(1 for e, s in zip(education, salary) if e=='0' and s>0)

but, for the sake of discussion, let's say that the length of the stream is really large and I don't want an unneeded gigantic list and that the stream is not replicable (not really my example…).

    Is it possible to sum and count the stream at the same time?

7 Answers

You can filter it with a generator expression instead of turning it into a list; all you need to do is use () instead of [] -- this will process it in a "streaming" fashion instead of creating the whole thing in memory:

low_edu = (s for e, s in zip(education, salary) if e=='0' and s>0)

Then just add up the total and count together:

total, count = 0, 0
for salary in low_edu:
    total += salary
    count += 1

You could mash this into another generator expression with functools.reduce but it reads a lot cleaner with a simple loop imo.

I'd just use a simple loop without the generator expression.

total = count = 0
for e, s in zip(education, salary):
    if e == '0' and s > 0:
        total += s
        count += 1
mean = total / count

Though if you already only have the iterator of preprocessed values, then tzaman's would be my choice. Then again, you're considering to "count again my stream", so apparently you are in control of the code so could do it without the generator.

Measuring memory usage and runtime of the various working solutions for larger input (100,000 elements):

            original  406.72 KB   7.0 ms  result = 24.990380884305115
         Kelly_Bundy    0.24 KB  10.1 ms  result = 24.990380884305115
              tzaman    0.40 KB  11.6 ms  result = 24.990380884305115
         user2390182 6155.90 KB  16.0 ms  result = 24.990380884305115
 alparslan_mimaroğlu  441.23 KB  21.4 ms  result = 24.990380884305115
         user1740577  406.86 KB   7.2 ms  result = 24.990380884305115

Code (Try it online!):

from random import choices
from itertools import tee
import tracemalloc
import functools
from timeit import repeat

def original(education, salary):
    list_of_les = [s for e, s in zip(education, salary) if e=='0' and s>0]
    return sum(list_of_les)/len(list_of_les)

def Kelly_Bundy(education, salary):
    total = count = 0
    for e, s in zip(education, salary):
        if e == '0' and s > 0:
            total += s
            count += 1
    return total / count

def tzaman(education, salary):
    low_edu = (s for e, s in zip(education, salary) if e=='0' and s>0)
    total, count = 0, 0
    for salary in low_edu:
        total += salary
        count += 1
    return total / count

def user2390182(education, salary):
    s, c = map(sum, zip(*((s, 1) for e, s in zip(education, salary) if e=='0' and s>0)))
    return s / c

def alparslan_mimaroğlu(education, salary):
    knowns = map(lambda x: x[1] , filter(lambda x: x[0] == '0' and x[1] != 0, zip(education, salary)))
    (knowns1, knowns2) =tee(knowns, 2)
    return sum(knowns1) / sum(1 for _ in knowns2)

def user1740577(education, salary):
    sum_len = lambda t : (sum(t), len(t))
    s,c = sum_len([s for e, s in zip(education, salary) if e=='0' and s>0])
    return s / c

def main():
    n = 10 ** 5
    education = choices(['0', '1'], k=n)
    salary = choices(range(50), k=n)

    def test(function):
        tracemalloc.start()
        result = function(iter(education), iter(salary))
        memory = tracemalloc.get_traced_memory()[1]
        tracemalloc.stop()
        time = min(repeat(lambda: function(iter(education), iter(salary)), number=1))
        print('%20s' % function.__name__,
              '%7.2f KB ' % (memory / 1e3),
              '%4.1f ms ' % (time * 1e3),
              'result =', result)

    for _ in range(3):
        for function in original, Kelly_Bundy, tzaman, user2390182, alparslan_mimaroğlu, user1740577:
            test(function)
        print()

main()

Using reduce you can basically do arbitrary aggregations, even multiples at once. The following calculates the sum and count at the same time:

import functools
data = [1, 3, 5, 6, 2]
s, c = functools.reduce(lambda a, b: [a[0] + b, a[1] + 1], data, [0,0]) 
# outputs 17, 5

Similar to @luk2302's approach, but not using reduce, but simple back-n-forth transpositioning:

s, c = map(sum, zip(*((s, 1) for e, s in zip(education, salary) if e=='0' and s>0)))

Four memory-efficient ways if you do already have a preprocessed "stream", not the two sources:

Walrus:

def mean(stream):
    count = 0
    total = sum(x for x in stream if (count := count + 1))
    return total / count

Advancing an itertools.count in parallel:

def mean(stream):
    ctr = count()
    total = sum(x for x, _ in zip(stream, ctr))
    return total / next(ctr)

Or with an itemgetter:

def mean(stream):
    ctr = count()
    total = sum(map(itemgetter(0), zip(stream, ctr)))
    return total / next(ctr)

With enumerate:

def mean(stream):
    total = 0
    for count, x in enumerate(stream, 1):
        total += x
    return total / count

Demo:

from itertools import count
from operator import itemgetter

def mean1(stream):
    count = 0
    total = sum(x for x in stream if (count := count + 1))
    return total / count

def mean2(stream):
    ctr = count()
    total = sum(x for x, _ in zip(stream, ctr))
    return total / next(ctr)

def mean3(stream):
    ctr = count()
    total = sum(map(itemgetter(0), zip(stream, ctr)))
    return total / next(ctr)

def mean4(stream):
    total = 0
    for count, x in enumerate(stream, 1):
        total += x
    return total / count

for func in mean1, mean2, mean3, mean4:
    education = '0 1 0 0 0'.split()
    salary = [int(s) for s in '50 120 0 40 60'.split()] # 0 for missing data
    stream = (s for e, s in zip(education, salary) if e=='0' and s>0)
    print(func(stream))

Output:

50.0
50.0
50.0
50.0

Here's how Python itself does it, for its statistics.fmean. See the part Handle iterators that do not define __len__():

def fmean(data):
    """Convert data to floats and compute the arithmetic mean.
    This runs faster than the mean() function and it always returns a float.
    If the input dataset is empty, it raises a StatisticsError.
    >>> fmean([3.5, 4.0, 5.25])
    4.25
    """
    try:
        n = len(data)
    except TypeError:
        # Handle iterators that do not define __len__().
        n = 0
        def count(iterable):
            nonlocal n
            for n, x in enumerate(iterable, start=1):
                yield x
        total = fsum(count(data))
    else:
        total = fsum(data)
    try:
        return total / n
    except ZeroDivisionError:
        raise StatisticsError('fmean requires at least one data point') from None

I btw measured the memory usage of statistics.fmean and it took very little. So in case you have a huge stream and just want the mean, this function might be good.

More direct comparison of more solutions, giving each a stream of 100,000 numbers without the backstory data, and somewhat unifying their code for better comparison.

           with_list   900.16 KB   1.0 ms  result = 49.66821
          with_tuple   800.10 KB   1.1 ms  result = 49.66821
           with_loop     0.08 KB   8.1 ms  result = 49.66821
      with_transpose 12688.18 KB  22.7 ms  result = 49.66821
            with_tee   898.95 KB   5.5 ms  result = 49.66821
         with_reduce     0.30 KB  21.8 ms  result = 49.66821
       with_reduce_2     0.25 KB  19.9 ms  result = 49.66821
              walrus     0.29 KB   8.4 ms  result = 49.66821
   itertools_count_1     0.36 KB   7.3 ms  result = 49.66821
   itertools_count_2     0.27 KB   6.6 ms  result = 49.66821
    with_enumerate_1     0.15 KB   7.6 ms  result = 49.66821
    with_enumerate_2     0.45 KB   6.9 ms  result = 49.66821
     statistics_mean   900.87 KB  50.8 ms  result = 49.66821
    statistics_fmean     0.64 KB   7.8 ms  result = 49.66821

Code (Try it online!):

from random import choices
from itertools import tee, count
import tracemalloc
from functools import reduce
from timeit import repeat
import statistics
from operator import itemgetter

def with_list(stream):
    return sum(lst := list(stream)) / len(lst)

def with_tuple(stream):
    return sum(tpl := tuple(stream)) / len(tpl)

def with_loop(stream):
    total = count = 0
    for x in stream:
        total += x
        count += 1
    return total / count

def with_transpose(stream):
    total, count = map(sum, zip(*((x, 1) for x in stream)))
    return total / count

def with_tee(stream):
    it1, it2 = tee(stream)
    return sum(it1) / sum(1 for _ in it2)

def with_reduce(stream):
    total, count = reduce(lambda tc, x: [tc[0] + x, tc[1] + 1], stream, [0, 0])
    return total / count

def with_reduce_2(stream):
    total, count = reduce(lambda tc, x: (tc[0] + x, tc[1] + 1), stream, (0, 0))
    return total / count

def walrus(stream):
    count = 0
    total = sum(x for x in stream if (count := count + 1))
    return total / count

def itertools_count_1(stream):
    ctr = count()
    total = sum(x for x, _ in zip(stream, ctr))
    return total / next(ctr)

def itertools_count_2(stream):
    ctr = count()
    total = sum(map(itemgetter(0), zip(stream, ctr)))
    return total / next(ctr)

def with_enumerate_1(stream):
    total = 0
    for count, x in enumerate(stream, 1):
        total += x
    return total / count

def with_enumerate_2(stream):
    n = 0
    def count():
        nonlocal n
        for n, x in enumerate(stream, start=1):
            yield x
    return sum(count()) / n

def statistics_mean(stream):
    return statistics.mean(stream)

def statistics_fmean(stream):
    return statistics.fmean(stream)

functions = [
    with_list,
    with_tuple,
    with_loop,
    with_transpose,
    with_tee,
    with_reduce,
    with_reduce_2,
    walrus,
    itertools_count_1,
    itertools_count_2,
    with_enumerate_1,
    with_enumerate_2,
    statistics_mean,
    statistics_fmean,
]

def main():
    n = 10 ** 5
    numbers = choices(range(100), k=n)

    def test(function):
        stream = iter(numbers)
        tracemalloc.start()
        result = function(stream)
        memory = tracemalloc.get_traced_memory()[1]

        tracemalloc.stop()
        time = min(repeat(lambda: function(iter(numbers)), number=1, repeat=30))

        print('%20s' % function.__name__,
              '%8.2f KB ' % (memory / 1e3),
              '%4.1f ms ' % (time * 1e3),
              'result =', result)

    for _ in range(3):
        for function in functions:
            test(function)
        print()

main()
Related