Neat redistribution of association values

If I understood correctly the goal, then may be you can benefit from generalizing this a bit.

The following function will reverse your redistribution rules, taking two functions - combiner and reducer:

ClearAll[reverse]
reverse[assoc_Association?AssociationQ, reducer_, combiner_] := 
  Merge[
    Flatten @ Map[reverse[#, reducer]&, Normal @ assoc], 
    combiner
  ]

reverse[key_ -> value_Association?AssociationQ, reducer_]:=
  Map[
    reverse[key, #, reducer]&, 
    Normal @ value
  ]

reverse[key_, other_ -> weight_, reducer_]:= other -> reducer[key, weight]

Now, you get what you want with:

reverse[
  distribution,
  Times[#2, original[#1]] &, 
  Total
]

(* <|"A" -> {0.5 a1, 0.5 a2}, "B" -> {0.5 a1, 0.5 a2}, "C" -> {b1 + c1, b2 + c2}|> *)

The code is somewhat more verbose than in your solution, but I think it is also somewhat more understandable.


We can write it without the need for variable mark as follows.

Merge[ KeyValueMap[#1*#2 &, distribution], Total ] /. original

<|"A" -> 0.5 {a1, a2}, "B" -> 0.5 {a1, a2}, "C" -> {b1, b2} + {c1, c2}|>

More so, the alternative solution provided by @Kuba using operator form is much more readable I believe.

ReplaceAll[original] @ Merge[Total] @ KeyValueMap[Times] @ distribution

<|"A" -> 0.5 {a1, a2}, "B" -> 0.5 {a1, a2}, "C" -> {b1, b2} + {c1, c2}|>


I don't think this answer is so much more concise than the implementation in the question. Far from it. I present it here mostly to hint at a different implementation for distribution, if such a reformulation makes sense in the context of the question.

Please consider redefining distribution to the definition below:

distribution2 = <|
  "A" -> <|"a" -> .5, "b" -> 0., "c" -> 0.|>,
  "B" -> <|"a" -> .5, "b" -> 0., "c" -> 0.|>,
  "C" -> <|"a" -> 0., "b" -> 1., "c" -> 1.|>
 |>;

Then, it's just a matter of 'applying' the distribution in a sense:

Map[(KeyValueMap[#2 original[#1] &, #] &) /* Total, distribution2]

evaluates to

<|"A" -> {0. + 0.5 a1, 0. + 0.5 a2}, 
  "B" -> {0. + 0.5 a1, 0. + 0.5 a2}, 
  "C" -> {0. + 1. b1 + 1. c1, 0. + 1. b2 + 1. c2}|>

You can add Chop if you need better looking results:

Map[(KeyValueMap[#2 original[#1] &, #] &) /* Total /* Chop, distribution2]
<|"A" -> {0.5 a1, 0.5 a2}, 
  "B" -> {0.5 a1, 0.5 a2}, 
  "C" -> {1. b1 + 1. c1, 1. b2 + 1. c2}|>