Weighted average of all values given 2d lists and library of values

Question

I have three different lists:

a = [[7,5,8],[3,4]]

value = [9, 5, 7, 8, 3, 4, 10]
weight = [100.0, 65.0, 25.0, 25.0, 100.0, 65.0, 25.0]
capacity = [1000.0, 15.0, 700.0, 700.0, 1000.0, 15.0, 700.0]

I am trying to find the weight and capacity of the sum of all the values in a given entry of a so for example in a[0]:

weight_a[0] = (700*25 + 15*65 + 700*25) / (700 + 15 + 700)
capacity_a[0] = 700 + 15 + 700

weight_a[1] = (1000*100 + 15*65) / (1000 + 15)
capacity_a[1] = 1000 + 15

So weight_a and capacity_a are the sum of the weighted_average and capacity of the entries in each list of a.

Value 7 has a weight of 25 and capacity of 700:

7 is value[2] so capacity[2] = 700 and weight[2] = 25

The weight_a equation is sum (weight of value * capacity of value) / sum of capacity of all values

capacity_a equation is the sum of capacity of all values

I have been stuck trying to formulate the problem and can't seem to figure it out. I am trying to zip(values, capacity, weight) but I am not sure how to access that given the 2d list a.

Any help is appreciated!

Answer 1

You can use map with a custom function, then use zip to unpack results:

from operator import itemgetter, mul

def calculator(k):
    k_set = set(k)
    getter = itemgetter(*(idx for idx, i in enumerate(value) if i in k_set))
    weights, caps = map(getter, (weight, capacity))
    capsum = sum(caps)
    return sum(map(mul, weights, caps)) / capsum, capsum

weight_a, capacity_a = zip(*map(calculator, a))

print(weight_a, capacity_a)

(25.424028268551236, 99.48275862068965) (1415.0, 1015.0)