Updated normalize to support dict

reachtarunhere · reachtarunhere · commit bbe9f3d17b5e · 2016-08-06T20:31:21.000+05:30
diff --git a/utils.py b/utils.py
@@ -228,10 +228,16 @@ def num_or_str(x):
             return str(x).strip()
 
 
-def normalize(numbers):
+def normalize(dist):
     """Multiply each number by a constant such that the sum is 1.0"""
-    total = float(sum(numbers))
-    return [(n / total) for n in numbers]
+    if isinstance(dist, dict):
+        total = sum(dist.values())
+        for key in dist:
+            dist[key] = dist[key] / total
+            assert 0 <= dist[key] <= 1, "Probabilities must be between 0 and 1."
+        return dist
+    total = sum(dist)
+    return [(n / total) for n in dist]
 
 
 def clip(x, lowest, highest):
