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 -class TopNContainer(object):
    14    """ maintains a sorted list of a particular number of data elements.
    15
    16    """
    17 -  def __init__(self,size,mostNeg=-1e99):
    18      """
    19      if size is negative, all entries will be kept in sorted order
    20      """
    21      self._size = size
    22      if(size>=0):
    23        self.best = [mostNeg]*self._size
    24        self.extras = [None]*self._size
    25      else:
    26        self.best = []
    27        self.extras = []
    28 -  def Insert(self,val,extra=None):
    29      """ only does the insertion if val fits """
    30      if self._size>=0:
    31        if val > self.best[0]:
    32          idx = bisect.bisect(self.best,val)
    33          # insert the new element
    34          if idx == self._size:
    35            self.best.append(val)
    36            self.extras.append(extra)
    37          else:
    38            self.best.insert(idx,val)
    39            self.extras.insert(idx,extra)
    40          # and pop off the head
    41          self.best.pop(0)
    42          self.extras.pop(0)
    43      else:
    44        idx = bisect.bisect(self.best,val)
    45        self.best.insert(idx,val)
    46        self.extras.insert(idx,extra)
    47
    48 -  def GetPts(self):
    49      """ returns our set of points """
    50      return self.best
    51 -  def GetExtras(self):
    52      """ returns our set of extras """
    53      return self.extras
    54
    55 -  def __len__(self):
    56      return self._size
    57 -  def __getitem__(self,which):
    58      return self.best[which],self.extras[which]
    59
    60 -  def reverse(self):
    61      self.best.reverse()
    62      self.extras.reverse()
    63
    64  if __name__ == '__main__':
    65    import random
    66    pts = [int(100*random.random()) for x in range(10)]
    67
    68    c = TopNContainer(4)
    69    for pt in pts:
    70      c.Insert(pt,extra=str(pt))
    71    print(c.GetPts())
    72    print(c.GetExtras())
    73