From 5d9567e8498704f8d693f69b868919df4a1bbda2 Mon Sep 17 00:00:00 2001
From: antononcube <antononcube@gmail.com>
Date: Mon, 9 Apr 2018 19:11:36 -0400
Subject: [PATCH] Implemented TrieClassify of Dataset and list of lists.

---
 TriesWithFrequencies.m | 28 ++++++++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/TriesWithFrequencies.m b/TriesWithFrequencies.m
index 45f07e6c..22758b59 100644
--- a/TriesWithFrequencies.m
+++ b/TriesWithFrequencies.m
@@ -573,12 +573,12 @@ Mathematica is (C) Copyright 1988-2016 Wolfram Research, Inc.
 Options[TrieClassify] := {"Default" -> None};
 
 TrieClassify[tr_, record_, opts : OptionsPattern[]] :=
-    TrieClassify[tr, record, "Decision", opts];
+    TrieClassify[tr, record, "Decision", opts] /; FreeQ[{opts}, "Probability"|"TopProbabilities"];
 
 TrieClassify[tr_, record_, "Decision", opts : OptionsPattern[]] :=
     First@Keys@TrieClassify[tr, record, "Probabilities", opts];
 
-TrieClassify[tr_, record_, "Probability" -> class_] :=
+TrieClassify[tr_, record_, "Probability" -> class_, opts : OptionsPattern[]] :=
     Lookup[TrieClassify[tr, record, "Probabilities"], class, 0];
 
 TrieClassify[tr_, record_, "TopProbabilities", opts : OptionsPattern[]] :=
@@ -596,6 +596,30 @@ Mathematica is (C) Copyright 1988-2016 Wolfram Research, Inc.
       ]
     ];
 
+TrieClassify[tr_, records:(_Dataset|{_List..}), "Decision", opts : OptionsPattern[]] :=
+    First @* Keys @* TakeLargest[1] /@ TrieClassify[tr, records, "Probabilities", opts];
+
+TrieClassify[tr_, records:(_Dataset|{_List..}), "Probability" -> class_, opts : OptionsPattern[]] :=
+    Map[Lookup[#, class, 0]&, TrieClassify[tr, records, "Probabilities"] ];
+
+TrieClassify[tr_, records:(_Dataset|{_List..}), "TopProbabilities", opts : OptionsPattern[]] :=
+    Map[ Select[#, # > 0 &]&, TrieClassify[tr, records, "Probabilities", opts] ];
+
+TrieClassify[tr_, records:(_Dataset|{_List..}), "TopProbabilities" -> n_Integer, opts : OptionsPattern[]] :=
+    Map[TakeLargest[#, UpTo[n]]&, TrieClassify[tr, records, "Probabilities", opts] ];
+
+TrieClassify[tr_, records:(_Dataset|{_List..}), "Probabilities", opts:OptionsPattern[] ] :=
+    Block[{clRes, classLabels, stencil},
+
+      clRes = Map[ TrieClassify[tr, #, "Probabilities", opts] &, Normal@records ];
+
+      classLabels = Union[Flatten[Normal[Keys /@ clRes]]];
+
+      stencil = AssociationThread[classLabels -> 0];
+
+      KeySort[Join[stencil, #]] & /@ clRes
+    ];
+
 End[]
 
 EndPackage[]
\ No newline at end of file