RandomTrackingDict now tracks properly

- it didn't put the correct keys at the end
This commit is contained in:
Peter Šurda 2018-02-01 14:43:14 +01:00
parent 290b87a49f
commit 167d946435
Signed by: PeterSurda
GPG Key ID: 0C5F50C0B5F37D87

View File

@ -31,21 +31,37 @@ class RandomTrackingDict(object):
self.dictionary[key] = [self.len, value] self.dictionary[key] = [self.len, value]
self.len += 1 self.len += 1
def _swap(self, i1, i2):
with self.lock:
key1 = self.indexDict[i1]
key2 = self.indexDict[i2]
self.indexDict[i1] = key2
self.indexDict[i2] = key1
self.dictionary[key1][0] = i2
self.dictionary[key2][0] = i1
# for quick reassignment
return i2
def __delitem__(self, key): def __delitem__(self, key):
if not key in self.dictionary: if not key in self.dictionary:
raise KeyError raise KeyError
with self.lock: with self.lock:
index = self.dictionary[key][0] index = self.dictionary[key][0]
self.indexDict[index] = self.indexDict[self.len - 1] # not pending
self.dictionary[self.indexDict[index]][0] = index if index < self.len - self.pendingLen:
# left of pending part
index = self._swap(index, self.len - self.pendingLen - 1)
# pending
else:
self.pendingLen -= 1
# end
self._swap(index, self.len - 1)
# if the following del is batched, performance of this single # if the following del is batched, performance of this single
# operation can improve 4x, but it's already very fast so we'll # operation can improve 4x, but it's already very fast so we'll
# ignore it for the time being # ignore it for the time being
del self.indexDict[-1] del self.indexDict[-1]
del self.dictionary[key] del self.dictionary[key]
self.len -= 1 self.len -= 1
if index >= self.len - self.pendingLen:
self.pendingLen -= 1
def setMaxPending(self, maxPending): def setMaxPending(self, maxPending):
self.maxPending = maxPending self.maxPending = maxPending
@ -63,16 +79,13 @@ class RandomTrackingDict(object):
available = self.len - self.pendingLen available = self.len - self.pendingLen
if count > available: if count > available:
count = available count = available
retval = random.sample(self.indexDict[:self.len - self.pendingLen], count) randomIndex = random.sample(range(self.len - self.pendingLen), count)
for i in retval[::-1]: retval = [self.indexDict[i] for i in randomIndex]
for i in sorted(randomIndex, reverse=True):
# swap with one below lowest pending # swap with one below lowest pending
self._swap(i, self.len - self.pendingLen - 1)
self.pendingLen += 1 self.pendingLen += 1
swapKey = self.indexDict[-self.pendingLen]
curIndex = self.dictionary[i][0]
self.indexDict[-self.pendingLen] = i
self.indexDict[curIndex] = swapKey
self.dictionary[i][0] = self.len - self.pendingLen
self.dictionary[swapKey][0] = curIndex
self.lastPoll = time() self.lastPoll = time()
return retval return retval