Improve documentation.

[pyutils.git] / src / pyutils / math_utils.py

diff --git a/src/pyutils/math_utils.py b/src/pyutils/math_utils.py
index 97bb635ac7af5cb31a680338a87ffaecbb4adce5..10a9fb774b6f4bea5810f439774184398a33ccb3 100644 (file)
--- a/src/pyutils/math_utils.py
+++ b/src/pyutils/math_utils.py
@@ -21,6 +21,8 @@ class NumericPopulation(object):
     >>> pop.add_number(1)
     >>> pop.add_number(10)
     >>> pop.add_number(3)
+    >>> len(pop)
+    3
     >>> pop.get_median()
     3
     >>> pop.add_number(7)
@@ -59,6 +61,15 @@ class NumericPopulation(object):
         if not self.minimum or number < self.minimum:
             self.minimum = number
 
+    def __len__(self):
+        """Return the population size."""
+        n = 0
+        if self.highers:
+            n += len(self.highers)
+        if self.lowers:
+            n += len(self.lowers)
+        return n
+
     def _rebalance(self):
         if len(self.lowers) - len(self.highers) > 1:
             heappush(self.highers, -heappop(self.lowers))
@@ -66,7 +77,7 @@ class NumericPopulation(object):
             heappush(self.lowers, -heappop(self.highers))
 
     def get_median(self) -> float:
-        """Returns the approximate median (p50) so far in O(1) time."""
+        """Returns the approximate median (p50) so far in :math:`O(1)` time."""
 
         if len(self.lowers) == len(self.highers):
             return -self.lowers[0]
@@ -76,14 +87,14 @@ class NumericPopulation(object):
             return self.highers[0]
 
     def get_mean(self) -> float:
-        """Returns the mean (arithmetic mean) so far in O(1) time."""
+        """Returns the mean (arithmetic mean) so far in :math:`O(1)` time."""
 
-        count = len(self.lowers) + len(self.highers)
+        count = len(self)
         return self.aggregate / count
 
     def get_mode(self) -> Tuple[float, int]:
         """Returns the mode (most common member in the population)
-        in O(n) time."""
+        in :math:`O(n)` time."""
 
         count: Dict[float, int] = collections.defaultdict(int)
         for n in self.lowers:
@@ -93,7 +104,7 @@ class NumericPopulation(object):
         return dict_utils.item_with_max_value(count)
 
     def get_stdev(self) -> float:
-        """Returns the stdev so far in O(n) time."""
+        """Returns the stdev so far in :math:`O(n)` time."""
 
         mean = self.get_mean()
         variance = 0.0
@@ -116,13 +127,13 @@ class NumericPopulation(object):
 
     def get_percentile(self, n: float) -> float:
         """Returns the number at approximately pn% (i.e. the nth percentile)
-        of the distribution in O(n log n) time.  Not thread-safe;
+        of the distribution in :math:`O(n log_2 n)` time.  Not thread-safe;
         does caching across multiple calls without an invocation to
         add_number for perf reasons.
         """
         if n == 50:
             return self.get_median()
-        count = len(self.lowers) + len(self.highers)
+        count = len(self)
         self._create_sorted_copy_if_needed(count)
         assert self.sorted_copy
         index = round(count * (n / 100.0))