|
5 | 5 | # Link: http://code.activestate.com/recipes/576930/
|
6 | 6 |
|
7 | 7 | # Cython version: Wes McKinney
|
8 |
| -from random import random |
9 |
| - |
10 |
| -from libc.math cimport log |
11 |
| - |
12 |
| -import numpy as np |
13 |
| - |
14 |
| - |
15 |
| -# MSVC does not have log2! |
16 |
| - |
17 |
| -cdef double Log2(double x): |
18 |
| - return log(x) / log(2.) |
19 |
| - |
20 |
| - |
21 |
| -# TODO: optimize this, make less messy |
22 |
| - |
23 |
| -cdef class Node: |
24 |
| - # cdef public: |
25 |
| - # double value |
26 |
| - # list next |
27 |
| - # list width |
28 |
| - |
29 |
| - def __init__(self, double value, list next, list width): |
30 |
| - self.value = value |
31 |
| - self.next = next |
32 |
| - self.width = width |
33 |
| - |
34 |
| - |
35 |
| -# Singleton terminator node |
36 |
| -NIL = Node(np.inf, [], []) |
37 |
| - |
38 |
| - |
39 |
| -cdef class IndexableSkiplist: |
40 |
| - """ |
41 |
| - Sorted collection supporting O(lg n) insertion, removal, and |
42 |
| - lookup by rank. |
43 |
| - """ |
44 |
| - # cdef: |
45 |
| - # Py_ssize_t size, maxlevels |
46 |
| - # Node head |
47 |
| - |
48 |
| - def __init__(self, expected_size=100): |
49 |
| - self.size = 0 |
50 |
| - self.maxlevels = int(1 + Log2(expected_size)) |
51 |
| - self.head = Node(np.NaN, [NIL] * self.maxlevels, [1] * self.maxlevels) |
52 |
| - |
53 |
| - def __len__(self): |
54 |
| - return self.size |
55 |
| - |
56 |
| - def __getitem__(self, i): |
57 |
| - return self.get(i) |
58 |
| - |
59 |
| - cpdef get(self, Py_ssize_t i): |
60 |
| - cdef: |
61 |
| - Py_ssize_t level |
62 |
| - Node node |
63 |
| - |
64 |
| - node = self.head |
65 |
| - i += 1 |
66 |
| - |
67 |
| - for level in range(self.maxlevels - 1, -1, -1): |
68 |
| - while node.width[level] <= i: |
69 |
| - i -= node.width[level] |
70 |
| - node = node.next[level] |
71 |
| - |
72 |
| - return node.value |
73 |
| - |
74 |
| - cpdef insert(self, double value): |
75 |
| - cdef: |
76 |
| - Py_ssize_t level, steps, d |
77 |
| - Node node, prevnode, newnode, next_at_level, tmp |
78 |
| - list chain, steps_at_level |
79 |
| - |
80 |
| - # find first node on each level where node.next[levels].value > value |
81 |
| - chain = [None] * self.maxlevels |
82 |
| - steps_at_level = [0] * self.maxlevels |
83 |
| - node = self.head |
84 |
| - |
85 |
| - for level in range(self.maxlevels - 1, -1, -1): |
86 |
| - next_at_level = node.next[level] |
87 |
| - |
88 |
| - while next_at_level.value <= value: |
89 |
| - steps_at_level[level] = (steps_at_level[level] + |
90 |
| - node.width[level]) |
91 |
| - node = next_at_level |
92 |
| - next_at_level = node.next[level] |
93 |
| - |
94 |
| - chain[level] = node |
95 |
| - |
96 |
| - # insert a link to the newnode at each level |
97 |
| - d = min(self.maxlevels, 1 - int(Log2(random()))) |
98 |
| - newnode = Node(value, [None] * d, [None] * d) |
99 |
| - steps = 0 |
100 |
| - |
101 |
| - for level in range(d): |
102 |
| - prevnode = chain[level] |
103 |
| - newnode.next[level] = prevnode.next[level] |
104 |
| - prevnode.next[level] = newnode |
105 |
| - newnode.width[level] = (prevnode.width[level] - steps) |
106 |
| - prevnode.width[level] = steps + 1 |
107 |
| - steps += steps_at_level[level] |
108 |
| - |
109 |
| - for level in range(d, self.maxlevels): |
110 |
| - (<Node>chain[level]).width[level] += 1 |
111 |
| - |
112 |
| - self.size += 1 |
113 |
| - |
114 |
| - cpdef remove(self, double value): |
115 |
| - cdef: |
116 |
| - Py_ssize_t level, d |
117 |
| - Node node, prevnode, tmpnode, next_at_level |
118 |
| - list chain |
119 |
| - |
120 |
| - # find first node on each level where node.next[levels].value >= value |
121 |
| - chain = [None] * self.maxlevels |
122 |
| - node = self.head |
123 |
| - |
124 |
| - for level in range(self.maxlevels - 1, -1, -1): |
125 |
| - next_at_level = node.next[level] |
126 |
| - while next_at_level.value < value: |
127 |
| - node = next_at_level |
128 |
| - next_at_level = node.next[level] |
129 |
| - |
130 |
| - chain[level] = node |
131 |
| - |
132 |
| - if value != (<Node>(<Node>(<Node>chain[0]).next)[0]).value: |
133 |
| - raise KeyError('Not Found') |
134 |
| - |
135 |
| - # remove one link at each level |
136 |
| - d = len((<Node>(<Node>(<Node>chain[0]).next)[0]).next) |
137 |
| - |
138 |
| - for level in range(d): |
139 |
| - prevnode = chain[level] |
140 |
| - tmpnode = prevnode.next[level] |
141 |
| - prevnode.width[level] += tmpnode.width[level] - 1 |
142 |
| - prevnode.next[level] = tmpnode.next[level] |
143 |
| - |
144 |
| - for level in range(d, self.maxlevels): |
145 |
| - tmpnode = chain[level] |
146 |
| - tmpnode.width[level] -= 1 |
147 |
| - |
148 |
| - self.size -= 1 |
0 commit comments