6
$\begingroup$

I am looking for a function that is similar to KeyMap in that it takes a function fun to transform each key. However, when several different old keys map to the same new key, then KeyMap will only keep the last value.

Example:

KeyMap[First, <|{a, 1} -> 1, {a, 2} -> 2, {b, 1} -> 3|>]
(* <|a -> 2, b -> 3|> *)

Instead, I want to provide a combiner function comb to combine values:

keyCombine[First, <|{a, 1} -> 1, {a, 2} -> 2, {b, 1} -> 3|>, Identity]
(* <|a -> {1, 2}, b -> {3} *)

keyCombine[First, <|{a, 1} -> 1, {a, 2} -> 2, {b, 1} -> 3|>, Total]
(* <|a -> 3, b -> 3 *)

Is there a built-in function for this? If not, what is the "best" implementation? Take "best" to refer either to performance, elegance or compactness.

One possible implementation:

keyCombine[fun_, asc_?AssociationQ, comb_ : Identity] := Merge[MapAt[fun, Normal[asc], {All, 1}], comb]

Benchmarking:

SeedRandom[42];
aa = AssociationThread[RandomInteger[1000000, 100000], RandomInteger[1000000, 100000]];

Mine:

keyCombine[Mod[#, 5] &, aa, Total] // AbsoluteTiming

(* {7.48073, <|4 -> 9451454209, 2 -> 9485726007, 
  3 -> 9480421781, 0 -> 9443541021, 1 -> 9545354067|>} *)

Kuba:

keyCombineBy[aa, Mod[#, 5] &, Total] // AbsoluteTiming
(* {0.191946, <|4 -> 9451454209, 2 -> 9485726007, 
  3 -> 9480421781, 0 -> 9443541021, 1 -> 9545354067|>} *)

keyCombineBy2[aa, Mod[#, 5] &, Total] // AbsoluteTiming
(* {7.37696, <|4 -> 9451454209, 2 -> 9485726007, 
  3 -> 9480421781, 0 -> 9443541021, 1 -> 9545354067|>} *)
$\endgroup$
5
$\begingroup$
keyCombineBy[assoc_?AssociationQ, by_, post_] := GroupBy[
    Normal@assoc, by@*First -> Last, post
]

keyCombineBy[<|{a, 1} -> 1, {a, 2} -> 2, {b, 1} -> 3|>, First, ff]
<|a -> ff[{1, 2}], b -> ff[{3}]|>

This minimal modification makes it slightly faster:

keyCombineBy3[assoc_?AssociationQ, by_, post_] := GroupBy[
        Normal@assoc, by@*First , post @* Values
    ]

Or, alternatively (slow):

 keyCombineBy2[assoc_?AssociationQ, by_, post_
 ] := Merge[post] @ KeyValueMap[by@#1 -> #2 &] @ assoc
$\endgroup$
  • $\begingroup$ Benchmark added in original post. $\endgroup$ – Szabolcs Apr 11 '17 at 9:14
  • $\begingroup$ @Szabolcs thanks, I was aiming in elegance but good to know the difference. $\endgroup$ – Kuba Apr 11 '17 at 9:15
4
$\begingroup$

You may use Query to improve performance.

ClearAll[keyCombine];
keyCombine[fun_, asc_?AssociationQ, comb_ : Identity] :=
 asc //
   Query[Normal /* GroupBy[fun@*Keys]] //
  Query[All, (Values@# &) /* comb]

Here Values@# & is used instead of Values to circumvent the syntax sugar that works against our intention in this case. Since we have a list of rules under each Key from the first Query then [All, Values] maps all the way down to the values of the list of keys. I think this is syntax sugar but it might be a bug (Ideas?). In any case rolling our own pure function escapes this and lets us place Values on the list of rules.

With

SeedRandom[42];
aa = AssociationThread[RandomInteger[1000000, 100000], RandomInteger[1000000, 100000]];

Then

keyCombine[Mod[#, 5] &, aa, Total] // AbsoluteTiming
{0.158329, 
   <|4 -> 9451454209, 2 -> 9485726007, 3 -> 9480421781, 0 -> 9443541021, 
     1 -> 9545354067|>}

Or, alternatively

ClearAll[keyCombine2];
keyCombine2[fun_, asc_?AssociationQ, comb_ : Identity] :=
 Query[Normal /* GroupBy[fun@*Keys] /* Map[comb@*Values]]@asc

This replaces both the second Query and the pure Values function with Map in the first Query.

Then

keyCombine2[Mod[#, 5] &, aa, Total] // AbsoluteTiming
{0.16708, 
   <|4 -> 9451454209, 2 -> 9485726007, 3 -> 9480421781, 0 -> 9443541021, 
     1 -> 9545354067|>}

This is ever so slightly slower than keyCombine above but some may find it easier to read.

Hope this helps.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.