rdkit.DataStructs.TopNContainer

Package rdkit :: Package DataStructs :: Module TopNContainer |
[hide private]
[frames] | no frames]

# Source Code for Module rdkit.DataStructs.TopNContainer

     1  # $Id$
     2  #
     3  #  Copyright (C) 2003-2013 Rational Discovery LLC
     4  #
     5  #   @@ All Rights Reserved @@
     6  #  This file is part of the RDKit.
     7  #  The contents are covered by the terms of the BSD license
     8  #  which is included in the file license.txt, found at the root
     9  #  of the RDKit source tree.
    10  #
    11  from __future__ import print_function
    12  import bisect
    13
    14
    15 -class TopNContainer(object):
    16    """ maintains a sorted list of a particular number of data elements.
    17
    18    """
    19
    20 -  def __init__(self, size, mostNeg=-1e99):
    21      """
    22      if size is negative, all entries will be kept in sorted order
    23      """
    24      self._size = size
    25      if (size >= 0):
    26        self.best = [mostNeg] * self._size
    27        self.extras = [None] * self._size
    28      else:
    29        self.best = []
    30        self.extras = []
    31
    32 -  def Insert(self, val, extra=None):
    33      """ only does the insertion if val fits """
    34      if self._size >= 0:
    35        if val > self.best[0]:
    36          idx = bisect.bisect(self.best, val)
    37          # insert the new element
    38          if idx == self._size:
    39            self.best.append(val)
    40            self.extras.append(extra)
    41          else:
    42            self.best.insert(idx, val)
    43            self.extras.insert(idx, extra)
    44          # and pop off the head
    45          self.best.pop(0)
    46          self.extras.pop(0)
    47      else:
    48        idx = bisect.bisect(self.best, val)
    49        self.best.insert(idx, val)
    50        self.extras.insert(idx, extra)
    51
    52 -  def GetPts(self):
    53      """ returns our set of points """
    54      return self.best
    55
    56 -  def GetExtras(self):
    57      """ returns our set of extras """
    58      return self.extras
    59
    60 -  def __len__(self):
    61      if self._size >= 0:
    62        return self._size
    63      else:
    64        return len(self.best)
    65
    66 -  def __getitem__(self, which):
    67      return self.best[which], self.extras[which]
    68
    69 -  def reverse(self):
    70      self.best.reverse()
    71      self.extras.reverse()
    72
    73
    74 -def _exampleCode():
    75    import random
    76    random.seed(0)
    77    pts = [int(100 * random.random()) for _ in range(10)]
    78
    79    c = TopNContainer(4)
    80    for pt in pts:
    81      c.Insert(pt, extra=str(pt))
    82    print(c.GetPts())
    83    print(c.GetExtras())
    84
    85
    86  if __name__ == '__main__':  # pragma: nocover
    87    _exampleCode()
    88