12
$\begingroup$

Especially in tree-structured data, keys and values, being two sides of the same coin, frequently interchange roles. The ability to aggregate by value but not by key upsets the balance.

How to improve on this workaround query - here defined only in operator form?

keyGroupBy[f_][expr_] :=  expr // AssociationMap[First[#] -> # &] // GroupBy[First /* f] // 
      Map[Values /* Association]

For example, given:

    data = Range[5] // AssociationMap[foo]

(* <|1 -> foo[1], 2 -> foo[2], 3 -> foo[3], 4 -> foo[4], 5 -> foo[5]|> *)

The desired output is:

data // keyGroupBy[PrimeQ]

(* <|False -> <|1 -> foo[1], 4 -> foo[4]|>, 
 True -> <|2 -> foo[2], 3 -> foo[3], 5 -> foo[5]|>|> *) 

A fundamental constraint is that keys are transparent to functions accessing Association. Further, Association supports only data with unique keys. For example, GroupBy nested in AssociationMap will drop all but the last key-value pair with a given key.

$\endgroup$
  • 1
    $\begingroup$ You can consider replacing AssociationMap[First[#] -> # &] by Normal. Here's my right-to-left version: Association /@ GroupBy[PrimeQ@*First] @ Normal[data] $\endgroup$ – Szabolcs Sep 9 '14 at 1:30
  • $\begingroup$ @Szabolcs, good, here's a version of your approach closer to my original expression: data // Normal // GroupBy[First /* PrimeQ] // Map[Association] $\endgroup$ – alancalvitti Sep 9 '14 at 1:39
  • $\begingroup$ @Szabolcs Why not post an answer? I can't find anything faster than that. $\endgroup$ – Mr.Wizard Sep 9 '14 at 4:31
  • $\begingroup$ @Mr.Wizard I wasn't going for speed at all and it's a minimal modification. Also, I'm not particularly confident in this area yet. $\endgroup$ – Szabolcs Sep 9 '14 at 4:49
13
$\begingroup$

Since Szabolcs apparently declined to post an answer here is one of my own. The best method appears to be the one that he proposed and I take no credit for it. I do however offer a couple of inferior alternatives that might hopefully inspire some other approach.

Code

(* Original function for reference *)
keyGroupBy[f_][expr_] := 
 expr // AssociationMap[First[#] -> # &] // GroupBy[First /* f] // 
  Map[Values /* Association]

(* Szabolcs's recommendation *)
keyGroupBy2[f_][expr_] := Association /@ GroupBy[Normal@expr, f@*Keys]

(* reverse, GroupBy, and reverse again *)
keyGroupBy3[f_][expr_] := AssociationMap[Reverse] // # /@ GroupBy[#@expr, f] & 

(* manual Sow and Reap *)
keyGroupBy4[f_][expr_] :=
 Reap[
   AssociationMap[Sow[#, {f@Keys@#}] &, expr];,
   _,
   # -> Association[#2] &
 ][[2]]

(* method from mfvonh *)
keyGroupBy5[f_][expr_] := AssociationMap[expr] /@ GroupBy[Keys[expr], f]

Timings

Needs["GeneralUtilities`"]

BenchmarkPlot[
 {keyGroupBy, keyGroupBy2, keyGroupBy3, keyGroupBy4, keyGroupBy5}[PrimeQ] // Through,
 foo ~AssociationMap~ Range[#] &,
 2^Range[30],
 TimeConstraint -> 15
]

enter image description here

Obviously keyGroupBy and keyGroupBy3 are slower than keyGroupBy2 in this particular test. keyGroupBy4 appears to have some promise with larger sets but additional testing indicates that it remains several times slower than keyGroupBy2.

| improve this answer | |
$\endgroup$
  • 1
    $\begingroup$ Another: keyGroupBy5[f_][expr_] := AssociationMap[expr] /@ GroupBy[Keys[expr], f] $\endgroup$ – mfvonh Sep 9 '14 at 8:05
  • $\begingroup$ @mfvonh Thanks! I included it in the timings. $\endgroup$ – Mr.Wizard Sep 9 '14 at 11:41
  • 1
    $\begingroup$ A simplification of my version, GroupBy[Normal[data], PrimeQ@*First, Association]. Also, thanks for posting! $\endgroup$ – Szabolcs Sep 9 '14 at 20:48
  • $\begingroup$ @Mr.Wizard, Thx for the analysis. Also Szabolcs x2, mfvonh, et al. $\endgroup$ – alancalvitti Sep 10 '14 at 17:10
6
$\begingroup$

The question of alternate implementations has already been answered nicely by Mr.Wizard.

I will however answer the question that is implicit in the title of the question: should KeyGroupBy exist in the language? I think the answer is clearly 'yes'.

| improve this answer | |
$\endgroup$
  • $\begingroup$ That's good to hear. $\endgroup$ – alancalvitti Sep 10 '14 at 17:18
  • $\begingroup$ KeyMapAt would also be useful. $\endgroup$ – alancalvitti Nov 14 '14 at 23:37

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.