CODEWITHSUNDEEP
pythonlistpython-2.7
Ibe
Ibe

Reputation: 6035

Sum and length of lists of list

List a has zero values which have non-zero values in list b at corresponding points:

a = [[0, 2, 2], [2, 0, 0], [2, 0, 1]]
b = [[1, 0, 0], [0, 2, 3], [0, 1, 0]]

I need sum of each list in b and length of non-zero values of each list in a. I get sum of lists in list b with:

s = [sum(i) for i in b] # output: [1, 5, 1]

But how do I get length of non-zero values in each list of a (expected output: l = [2, 1, 2]) and add resultant values (s/l) to each non-zero value of a?

Expected output:

[[0, 2.5, 2.5], [7, 0, 0], [2.5, 0, 1.5]]

Upvotes: 3

Views: 4662

Answers (1)

Martijn Pieters
Martijn Pieters

Reputation: 1122172

Counting non-zero values can be done with sum() too:

l = [sum(v != 0 for v in i) for i in a]

This uses a property of Python booleans: they are a subclass of int and have integer values 0 and 1 for False and True respectively. As a result, summing booleans gives you an integer result letting you count how many of the v != 0 tests are true.

To calculate s / l, zip their results together with the nested lists in a, then use a conditional expression to add the ratio to each non-zero value:

s = [sum(i) for i in b]
l = [sum(v != 0 for v in i) for i in a]
result = [[0 if not v else v + (float(sumb) / count) for v in i]
          for i, sumb, count in zip(a, s, l)]

Demo:

>>> a = [[0, 2, 2], [2, 0, 0], [2, 0, 1]]
>>> b = [[1, 0, 0], [0, 2, 3], [0, 1, 0]]
>>> s = [sum(i) for i in b]
>>> l = [sum(v != 0 for v in i) for i in a]
>>> l
[2, 1, 2]
>>> [[0 if not v else v + (float(sumb) / count) for v in i]
...  for i, sumb, count in zip(a, s, l)]
[[0, 2.5, 2.5], [7.0, 0, 0], [2.5, 0, 1.5]]

You could pre-calculate the s / l ratio first; this can shave of a few cycles if the nested lists in a are very large:

ratio_sl = [sum(ib, 0.0) / sum(v != 0 for v in ia)  for ia, ib in zip(a, b)]
result = [[0 if not v else v + r for v in i] for i, r in zip(a, ratio_sl)]

Resulting output:

>>> ratio_sl = [sum(ib, 0.0) / sum(v != 0 for v in ia)  for ia, ib in zip(a, b)]
>>> ratio_sl
[0.5, 5.0, 0.5]
>>> [[0 if not v else v + r for v in i] for i, r in zip(a, ratio_sl)]
[[0, 2.5, 2.5], [7.0, 0, 0], [2.5, 0, 1.5]]

Here I stared sum() with a floating point 0.0 to ensure that Python 2 doesn't use integer division when calculating the ratios.

Upvotes: 2

Related Questions