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