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}|>