rbtree-ruby 0.2.2 → 0.3.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
data/lib/rbtree.rb CHANGED
@@ -56,7 +56,7 @@ class RBTree
56
56
 
57
57
  # Returns the number of key-value pairs stored in the tree.
58
58
  # @return [Integer] the number of entries in the tree
59
- attr_reader :size
59
+ attr_reader :key_count
60
60
 
61
61
  # Creates a new RBTree from the given arguments.
62
62
  #
@@ -72,11 +72,15 @@ class RBTree
72
72
 
73
73
  # Initializes a new RBTree.
74
74
  #
75
- # The tree can be initialized empty or populated with initial data from a Hash or Array.
75
+ # The tree can be initialized empty or populated with initial data from a Hash, Array, or Enumerator.
76
+ # A block can also be provided to supply the initial data.
76
77
  #
77
78
  # @param args [Hash, Array, nil] optional initial data
79
+ # @param overwrite [Boolean] whether to overwrite existing keys (default: true)
80
+ # @yieldreturn [Object] optional initial data
78
81
  # - If a Hash is provided, each key-value pair is inserted into the tree
79
82
  # - If an Array is provided, it should contain [key, value] pairs
83
+ # - If a block is provided, it is yielded to get the source data
80
84
  # - If no arguments are provided, an empty tree is created
81
85
  # @raise [ArgumentError] if arguments are invalid
82
86
  # @example Create an empty tree
@@ -85,7 +89,9 @@ class RBTree
85
89
  # tree = RBTree.new({1 => 'one', 2 => 'two'})
86
90
  # @example Create from an array
87
91
  # tree = RBTree.new([[1, 'one'], [2, 'two']])
88
- def initialize(*args)
92
+ # @example Create with overwrite: false
93
+ # tree = RBTree.new([[1, 'one'], [1, 'uno']], overwrite: false)
94
+ def initialize(*args, overwrite: true, &block)
89
95
  @nil_node = Node.new
90
96
  @nil_node.color = Node::BLACK
91
97
  @nil_node.left = @nil_node
@@ -94,28 +100,30 @@ class RBTree
94
100
  @min_node = @nil_node
95
101
  @hash_index = {} # Hash index for O(1) key lookup
96
102
  @node_pool = [] # Memory pool for recycling nodes
97
- @size = 0
98
-
99
- if args.any?
100
- source = args.size == 1 ? args.first : args
101
- case source
102
- when Hash
103
- source.each { |k, v| insert(k, v) }
104
- when Array
105
- source.each do |arg|
106
- key, value = arg
107
- insert(key, value)
108
- end
109
- else
110
- raise ArgumentError, "Invalid arguments"
111
- end
103
+ @key_count = 0
104
+
105
+ if args.size > 0 || block_given?
106
+ insert(*args, overwrite: overwrite, &block)
112
107
  end
113
108
  end
114
109
 
115
110
  # Checks if the tree is empty.
116
111
  #
117
112
  # @return [Boolean] true if the tree contains no elements, false otherwise
118
- def empty? = @root == @nil_node
113
+ def empty? = @hash_index.empty?
114
+
115
+ # Returns the number of key-value pairs stored in the tree.
116
+ # @return [Integer] the number of entries in the tree
117
+ def size = @key_count
118
+ alias :value_count :size
119
+
120
+ # Returns the minimum key without removing it.
121
+ #
122
+ # @return [Object, nil] the minimum key, or nil if tree is empty
123
+ # @example
124
+ # tree = RBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
125
+ # tree.min_key # => 1
126
+ def min_key = min_node&.key
119
127
 
120
128
  # Returns the minimum key-value pair without removing it.
121
129
  #
@@ -123,7 +131,15 @@ class RBTree
123
131
  # @example
124
132
  # tree = RBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
125
133
  # tree.min # => [1, "one"]
126
- def min = find_min
134
+ def min = min_node&.pair
135
+
136
+ # Returns the maximum key without removing it.
137
+ #
138
+ # @return [Object, nil] the maximum key, or nil if tree is empty
139
+ # @example
140
+ # tree = RBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
141
+ # tree.max_key # => 3
142
+ def max_key = max_node&.key
127
143
 
128
144
  # Returns the maximum key-value pair without removing it.
129
145
  #
@@ -131,7 +147,23 @@ class RBTree
131
147
  # @example
132
148
  # tree = RBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
133
149
  # tree.max # => [3, "three"]
134
- def max = find_max
150
+ def max = max_node&.pair
151
+
152
+ # Returns the first key-value pair without removing it.
153
+ #
154
+ # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
155
+ # @example
156
+ # tree = RBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
157
+ # tree.first # => [1, "one"]
158
+ def first = min
159
+
160
+ # Returns the last key-value pair without removing it.
161
+ #
162
+ # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
163
+ # @example
164
+ # tree = RBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
165
+ # tree.last # => [3, "three"]
166
+ def last = max
135
167
 
136
168
  # Checks if the tree contains the given key.
137
169
  #
@@ -139,11 +171,10 @@ class RBTree
139
171
  # @return [Boolean] true if the key exists in the tree, false otherwise
140
172
  # @example
141
173
  # tree = RBTree.new({1 => 'one', 2 => 'two'})
142
- # tree.has_key?(1) # => true
143
- # tree.has_key?(3) # => false
144
- def has_key?(key)
145
- @hash_index.key?(key)
146
- end
174
+ # tree.key?(1) # => true
175
+ # tree.key?(3) # => false
176
+ def has_key?(key) = @hash_index.key?(key)
177
+ alias :key? :has_key?
147
178
 
148
179
  # Retrieves the value associated with the given key.
149
180
  #
@@ -154,10 +185,22 @@ class RBTree
154
185
  # tree.get(1) # => "one"
155
186
  # tree[2] # => "two"
156
187
  # tree[3] # => nil
157
- def get(key)
158
- @hash_index[key]&.value
159
- end
160
- alias_method :[], :get
188
+ def value(key) = @hash_index[key]&.value
189
+ alias :get :value
190
+ alias :[] :value
191
+
192
+ # Returns the key with the key closest to the given key.
193
+ #
194
+ # This method requires keys to be numeric or support subtraction and abs methods.
195
+ #
196
+ # @param key [Numeric] the target key
197
+ # @return [Object, nil] the key, or nil if tree is empty
198
+ # @example
199
+ # tree = RBTree.new({1 => 'one', 5 => 'five', 10 => 'ten'})
200
+ # tree.nearest_key(4) # => 5
201
+ # tree.nearest_key(7) # => 5
202
+ # tree.nearest_key(8) # => 10
203
+ def nearest_key(key) = ((n = find_nearest_node(key)) == @nil_node)? nil : n.key
161
204
 
162
205
  # Returns the key-value pair with the key closest to the given key.
163
206
  #
@@ -171,11 +214,21 @@ class RBTree
171
214
  # tree.nearest(4) # => [5, "five"]
172
215
  # tree.nearest(7) # => [5, "five"]
173
216
  # tree.nearest(8) # => [10, "ten"]
174
- def nearest(key)
175
- return nil unless key.respond_to?(:-)
176
- n = find_nearest_node(key)
177
- n == @nil_node ? nil : n.pair
178
- end
217
+ def nearest(key) = ((n = find_nearest_node(key)) == @nil_node)? nil : n.pair
218
+
219
+ # Returns the key with the largest key that is smaller than the given key.
220
+ #
221
+ # If the key exists in the tree, returns the predecessor (previous element).
222
+ # If the key does not exist, returns the largest key with key < given key.
223
+ #
224
+ # @param key [Object] the reference key
225
+ # @return [Object, nil] the key, or nil if no predecessor exists
226
+ # @example
227
+ # tree = RBTree.new({1 => 'one', 3 => 'three', 5 => 'five', 7 => 'seven'})
228
+ # tree.prev_key(5) # => 3
229
+ # tree.prev_key(4) # => 3 (4 does not exist)
230
+ # tree.prev_key(1) # => nil (no predecessor)
231
+ def prev_key(key) = ((n = find_predecessor_node(key)) == @nil_node)? nil : n.key
179
232
 
180
233
  # Returns the key-value pair with the largest key that is smaller than the given key.
181
234
  #
@@ -189,10 +242,21 @@ class RBTree
189
242
  # tree.prev(5) # => [3, "three"]
190
243
  # tree.prev(4) # => [3, "three"] (4 does not exist)
191
244
  # tree.prev(1) # => nil (no predecessor)
192
- def prev(key)
193
- n = find_predecessor_node(key)
194
- n == @nil_node ? nil : n.pair
195
- end
245
+ def prev(key) = ((n = find_predecessor_node(key)) == @nil_node)? nil : n.pair
246
+
247
+ # Returns the key with the smallest key that is larger than the given key.
248
+ #
249
+ # If the key exists in the tree, returns the successor (next element).
250
+ # If the key does not exist, returns the smallest key with key > given key.
251
+ #
252
+ # @param key [Object] the reference key
253
+ # @return [Object, nil] the key, or nil if no successor exists
254
+ # @example
255
+ # tree = RBTree.new({1 => 'one', 3 => 'three', 5 => 'five', 7 => 'seven'})
256
+ # tree.succ_key(5) # => 7
257
+ # tree.succ_key(4) # => 5 (4 does not exist)
258
+ # tree.succ_key(1) # => 3 (1 does not exist)
259
+ def succ_key(key) = ((n = find_successor_node(key)) == @nil_node)? nil : n.key
196
260
 
197
261
  # Returns the key-value pair with the smallest key that is larger than the given key.
198
262
  #
@@ -206,70 +270,71 @@ class RBTree
206
270
  # tree.succ(5) # => [7, "seven"]
207
271
  # tree.succ(4) # => [5, "five"] (4 does not exist)
208
272
  # tree.succ(7) # => nil (no successor)
209
- def succ(key)
210
- n = find_successor_node(key)
211
- n == @nil_node ? nil : n.pair
212
- end
273
+ def succ(key) = ((n = find_successor_node(key)) == @nil_node)? nil : n.pair
213
274
 
214
- # Inserts or updates a key-value pair in the tree.
275
+ # Inserts one or more key-value pairs into the tree.
276
+ #
277
+ # This method supports both single entry insertion and bulk insertion.
278
+ #
279
+ # Single insertion:
280
+ # insert(key, value, overwrite: true)
281
+ #
282
+ # Bulk insertion:
283
+ # insert(hash, overwrite: true)
284
+ # insert(array_of_pairs, overwrite: true)
285
+ # insert(enumerator, overwrite: true)
286
+ # insert { data_source }
215
287
  #
216
288
  # If the key already exists and overwrite is true (default), the value is updated.
217
289
  # If overwrite is false and the key exists, the operation returns nil without modification.
218
290
  #
219
- # @param key [Object] the key to insert (must implement <=>)
220
- # @param value [Object] the value to associate with the key
291
+ # @param args [Object] key (and value) or source object
221
292
  # @param overwrite [Boolean] whether to overwrite existing keys (default: true)
222
- # @return [Boolean, nil] true if inserted/updated, nil if key exists and overwrite is false
223
- # @example
224
- # tree = RBTree.new
225
- # tree.insert(1, 'one') # => true
226
- # tree.insert(1, 'ONE') # => true (overwrites)
227
- # tree.insert(1, 'uno', overwrite: false) # => nil (no change)
228
- # tree[2] = 'two' # using alias
229
- def insert(key, value, overwrite: true)
230
- if (node = @hash_index[key])
231
- return nil unless overwrite
232
- node.value = value
233
- return true
234
- end
235
- y = @nil_node
236
- x = @root
237
- while x != @nil_node
238
- y = x
239
- cmp = key <=> x.key
240
- if cmp == 0
241
- return nil unless overwrite
242
- x.value = value
243
- return true
244
- elsif cmp < 0
245
- x = x.left
293
+ # @yieldreturn [Object] data source for bulk insertion
294
+ # @return [Boolean, nil] true if inserted/updated, nil if key exists and overwrite is false (for single insert)
295
+ # @example Single insert
296
+ # tree.insert(1, 'one')
297
+ # @example Bulk insert from Hash
298
+ # tree.insert({1 => 'one', 2 => 'two'})
299
+ # @example Bulk insert from Array
300
+ # tree.insert([[1, 'one'], [2, 'two']])
301
+ def insert(*args, overwrite: true, &block)
302
+ if args.size == 2
303
+ key, value = args
304
+ insert_entry(key, value, overwrite: overwrite)
305
+ else
306
+ source = nil
307
+ if args.empty? && block_given?
308
+ source = yield
309
+ elsif args.size == 1
310
+ source = args[0]
311
+ elsif args.empty?
312
+ return # No-op
246
313
  else
247
- x = x.right
314
+ raise ArgumentError, "wrong number of arguments (given #{args.size}, expected 0..2)"
315
+ end
316
+
317
+ return if source.nil?
318
+
319
+ unless source.respond_to?(:each)
320
+ raise ArgumentError, "Source must be iterable"
321
+ end
322
+
323
+ source.each do |*pair|
324
+ key, value = nil, nil
325
+ if pair.size == 1 && pair[0].is_a?(Array)
326
+ key, value = pair[0]
327
+ raise ArgumentError, "Invalid pair size: #{pair[0].size} (expected 2)" unless pair[0].size == 2
328
+ elsif pair.size == 2
329
+ key, value = pair
330
+ else
331
+ raise ArgumentError, "Invalid pair format: #{pair.inspect}"
332
+ end
333
+ insert_entry(key, value, overwrite: overwrite)
248
334
  end
249
335
  end
250
- z = allocate_node(key, value, Node::RED, @nil_node, @nil_node, @nil_node)
251
- z.parent = y
252
- if y == @nil_node
253
- @root = z
254
- elsif (key <=> y.key) < 0
255
- y.left = z
256
- else
257
- y.right = z
258
- end
259
- z.left = @nil_node
260
- z.right = @nil_node
261
- z.color = Node::RED
262
- insert_fixup(z)
263
- @size += 1
264
-
265
- if @min_node == @nil_node || (key <=> @min_node.key) < 0
266
- @min_node = z
267
- end
268
-
269
- @hash_index[key] = z # Add to hash index
270
- true
271
336
  end
272
- alias_method :[]=, :insert
337
+ alias :[]= :insert
273
338
 
274
339
  # Deletes the key-value pair with the specified key.
275
340
  #
@@ -279,12 +344,12 @@ class RBTree
279
344
  # tree = RBTree.new({1 => 'one', 2 => 'two'})
280
345
  # tree.delete(1) # => "one"
281
346
  # tree.delete(3) # => nil
282
- def delete(key)
283
- value = delete_node(key)
284
- return nil unless value
285
- @size -= 1
347
+ def delete_key(key)
348
+ return nil unless (value = (z = @hash_index[key])&.value)
349
+ delete_indexed_node(key)
286
350
  value
287
351
  end
352
+ alias :delete :delete_key
288
353
 
289
354
  # Removes and returns the minimum key-value pair.
290
355
  #
@@ -294,10 +359,10 @@ class RBTree
294
359
  # tree.shift # => [1, "one"]
295
360
  # tree.shift # => [2, "two"]
296
361
  def shift
297
- return nil if @min_node == @nil_node
298
- result = [@min_node.key, @min_node.value]
299
- delete(@min_node.key)
300
- result
362
+ return nil unless (n = @min_node) != @nil_node
363
+ pair = n.pair
364
+ delete(n.key)
365
+ pair
301
366
  end
302
367
 
303
368
  # Removes and returns the maximum key-value pair.
@@ -308,21 +373,46 @@ class RBTree
308
373
  # tree.pop # => [3, "three"]
309
374
  # tree.pop # => [2, "two"]
310
375
  def pop
311
- n = rightmost(@root)
312
- return nil if n == @nil_node
313
- result = n.pair
376
+ return nil unless (n = rightmost(@root)) != @nil_node
377
+ pair = n.pair
314
378
  delete(n.key)
315
- result
379
+ pair
316
380
  end
317
381
 
318
382
  # Removes all key-value pairs from the tree.
319
383
  #
320
384
  # @return [RBTree] self
321
385
  def clear
322
- @root = @nil_node
323
- @min_node = @nil_node
386
+ @root = @min_node = @nil_node
324
387
  @hash_index.clear
325
- @size = 0
388
+ @key_count = 0
389
+ self
390
+ end
391
+
392
+ # Iterates over all keys in ascending (or descending) order.
393
+ #
394
+ # @param reverse [Boolean] if true, iterate in descending order (default: false)
395
+ # @param safe [Boolean] if true, safe for modifications during iteration (default: false)
396
+ # @yield [key] each key in the tree
397
+ # @return [Enumerator, RBTree] an Enumerator if no block is given, self otherwise
398
+ # @example
399
+ # tree = RBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
400
+ # tree.keys { |k| puts k }
401
+ # # Output:
402
+ # # 1
403
+ # # 2
404
+ # # 3
405
+ #
406
+ # # Reverse iteration
407
+ # tree.keys(reverse: true) { |k| ... }
408
+ #
409
+ # # Safe iteration for modifications
410
+ # tree.keys(safe: true) do |k|
411
+ # tree.delete(k) if k.even?
412
+ # end
413
+ def keys(reverse: false, safe: false, &block)
414
+ return enum_for(:keys, reverse: reverse, safe: safe) { @key_count } unless block_given?
415
+ each(reverse: reverse, safe: safe) { |key, _| yield key }
326
416
  self
327
417
  end
328
418
 
@@ -348,7 +438,7 @@ class RBTree
348
438
  # tree.delete(k) if k.even?
349
439
  # end
350
440
  def each(reverse: false, safe: false, &block)
351
- return enum_for(:each, reverse: reverse, safe: safe) unless block_given?
441
+ return enum_for(:each, reverse: reverse, safe: safe) { size } unless block_given?
352
442
  if reverse
353
443
  traverse_all_desc(@root, safe: safe, &block)
354
444
  else
@@ -376,7 +466,10 @@ class RBTree
376
466
  # @yield [key, value] each key-value pair in the tree
377
467
  # @return [Enumerator, RBTree] an Enumerator if no block is given, self otherwise
378
468
  # @see #each
379
- def reverse_each(safe: false, &block) = each(reverse: true, safe: safe, &block)
469
+ def reverse_each(safe: false, &block)
470
+ return enum_for(:reverse_each, safe: safe) { size } unless block_given?
471
+ each(reverse: true, safe: safe, &block)
472
+ end
380
473
 
381
474
  # Retrieves all key-value pairs with keys less than the specified key.
382
475
  #
@@ -493,13 +586,9 @@ class RBTree
493
586
  #
494
587
  # @return [String] a human-readable representation of the tree
495
588
  def inspect
496
- if @size > 0
497
- content = first(5).map { |k, v| "#{k.inspect}=>#{v.inspect}" }.join(", ")
498
- suffix = @size > 5 ? ", ..." : ""
499
- "#<#{self.class}:0x#{object_id.to_s(16)} size=#{@size} {#{content}#{suffix}}>"
500
- else
501
- super
502
- end
589
+ content = take(5).map { |k, v| "#{k.inspect}=>#{v.inspect}" }.join(", ")
590
+ suffix = size > 5 ? ", ..." : ""
591
+ "#<#{self.class}:0x#{object_id.to_s(16)} size=#{size} {#{content}#{suffix}}>"
503
592
  end
504
593
 
505
594
  # Validates the red-black tree properties.
@@ -521,6 +610,59 @@ class RBTree
521
610
  # @!visibility private
522
611
  private
523
612
 
613
+ def min_node = ((n = @min_node) == @nil_node) ? nil : n
614
+
615
+ def max_node = ((n = rightmost(@root)) == @nil_node) ? nil : n
616
+
617
+ # Inserts a single key-value pair.
618
+ #
619
+ # @param key [Object] the key to insert
620
+ # @param value [Object] the value to associate with the key
621
+ # @param overwrite [Boolean] whether to overwrite existing keys (default: true)
622
+ # @return [Boolean, nil] true if inserted/updated, nil if key exists and overwrite is false
623
+ def insert_entry(key, value, overwrite: true)
624
+ if (node = @hash_index[key])
625
+ return nil unless overwrite
626
+ node.value = value
627
+ return true
628
+ end
629
+ y = @nil_node
630
+ x = @root
631
+ while x != @nil_node
632
+ y = x
633
+ cmp = key <=> x.key
634
+ if cmp == 0
635
+ return nil unless overwrite
636
+ x.value = value
637
+ return true
638
+ elsif cmp < 0
639
+ x = x.left
640
+ else
641
+ x = x.right
642
+ end
643
+ end
644
+ z = allocate_node(key, value, Node::RED, @nil_node, @nil_node, @nil_node)
645
+ z.parent = y
646
+ if y == @nil_node
647
+ @root = z
648
+ elsif (key <=> y.key) < 0
649
+ y.left = z
650
+ else
651
+ y.right = z
652
+ end
653
+ z.left = @nil_node
654
+ z.right = @nil_node
655
+ z.color = Node::RED
656
+ insert_fixup(z)
657
+
658
+ if @min_node == @nil_node || (key <=> @min_node.key) < 0
659
+ @min_node = z
660
+ end
661
+
662
+ @hash_index[key] = z # Add to hash index
663
+ true
664
+ end
665
+
524
666
  # Traverses the tree in ascending order (in-order traversal).
525
667
  #
526
668
  # @param node [Node] the current node
@@ -797,11 +939,7 @@ class RBTree
797
939
  #
798
940
  # @param key [Object] the key to delete
799
941
  # @return [Object, nil] the value of the deleted node, or nil if not found
800
- def delete_node(key)
801
- z = @hash_index.delete(key) # O(1) lookup and remove from index
802
- return nil unless z
803
- remove_node(z)
804
- end
942
+ def delete_indexed_node(key) = (z = @hash_index.delete(key)) && delete_node(z)
805
943
 
806
944
  # Removes a node from the tree and restores red-black properties.
807
945
  #
@@ -812,7 +950,7 @@ class RBTree
812
950
  #
813
951
  # @param z [Node] the node to remove
814
952
  # @return [Object] the value of the removed node
815
- def remove_node(z)
953
+ def delete_node(z)
816
954
  next_min_node = nil
817
955
  if z == @min_node
818
956
  if z.right != @nil_node
@@ -937,25 +1075,6 @@ class RBTree
937
1075
  v.parent = u.parent
938
1076
  end
939
1077
 
940
- # Searches for a node with the given key.
941
- #
942
- # @param key [Object] the key to search for
943
- # @return [Node] the found node, or @nil_node if not found
944
- def find_node(key)
945
- current = @root
946
- while current != @nil_node
947
- cmp = key <=> current.key
948
- if cmp == 0
949
- return current
950
- elsif cmp < 0
951
- current = current.left
952
- else
953
- current = current.right
954
- end
955
- end
956
- @nil_node
957
- end
958
-
959
1078
  # Finds the node with the closest key to the given key.
960
1079
  #
961
1080
  # Uses numeric distance (absolute difference) to determine proximity.
@@ -1110,7 +1229,7 @@ class RBTree
1110
1229
  # Returns the minimum key-value pair.
1111
1230
  #
1112
1231
  # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
1113
- def find_min = @min_node == @nil_node ? nil : @min_node.pair
1232
+ def find_min = ((n = @min_node) != @nil_node) && n.pair
1114
1233
 
1115
1234
  # Finds the rightmost (maximum) node in a subtree.
1116
1235
  #
@@ -1126,7 +1245,7 @@ class RBTree
1126
1245
  # Returns the maximum key-value pair.
1127
1246
  #
1128
1247
  # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
1129
- def find_max = (n = rightmost(@root)) == @nil_node ? nil : n.pair
1248
+ def find_max = ((n = rightmost(@root)) != @nil_node) && n.pair
1130
1249
 
1131
1250
  # Performs a left rotation on the given node.
1132
1251
  #
@@ -1189,7 +1308,8 @@ class RBTree
1189
1308
  # Allocates a new node or recycles one from the pool.
1190
1309
  # @return [Node]
1191
1310
  def allocate_node(key, value, color, left, right, parent)
1192
- if (node = @node_pool.pop)
1311
+ node = @node_pool.pop
1312
+ if node
1193
1313
  node.key = key
1194
1314
  node.value = value
1195
1315
  node.color = color
@@ -1198,8 +1318,10 @@ class RBTree
1198
1318
  node.parent = parent
1199
1319
  node
1200
1320
  else
1201
- Node.new(key, value, color, left, right, parent)
1321
+ node = Node.new(key, value, color, left, right, parent)
1202
1322
  end
1323
+ @key_count += 1
1324
+ node
1203
1325
  end
1204
1326
 
1205
1327
  # Releases a node back to the pool.
@@ -1210,6 +1332,7 @@ class RBTree
1210
1332
  node.parent = nil
1211
1333
  node.value = nil # Help GC
1212
1334
  @node_pool << node
1335
+ @key_count -= 1
1213
1336
  end
1214
1337
 
1215
1338
  # Recursively checks black height consistency.
@@ -1256,7 +1379,7 @@ end
1256
1379
  # == Features
1257
1380
  #
1258
1381
  # * Multiple values per key using arrays
1259
- # * Separate methods for single deletion (`delete_one`) vs. all deletions (`delete`)
1382
+ # * Separate methods for single deletion (`delete_value`) vs. all deletions (`delete_key`)
1260
1383
  # * Values for each key maintain insertion order
1261
1384
  # * Configurable access to first or last value via `:last` option
1262
1385
  #
@@ -1265,10 +1388,10 @@ end
1265
1388
  # For each key, values are stored in insertion order. Methods that access
1266
1389
  # a single value support a `:last` option to choose which end of the array:
1267
1390
  #
1268
- # * +get(key)+, +get_first(key)+ - returns first value (oldest)
1269
- # * +get(key, last: true)+, +get_last(key)+ - returns last value (newest)
1270
- # * +delete_one(key)+, +delete_first(key)+ - removes first value
1271
- # * +delete_one(key, last: true)+, +delete_last(key)+ - removes last value
1391
+ # * +get(key)+, +first_value(key)+ - returns first value (oldest)
1392
+ # * +get(key, last: true)+, +last_value(key)+ - returns last value (newest)
1393
+ # * +delete_value(key)+, +delete_first_value(key)+ - removes first value
1394
+ # * +delete_value(key, last: true)+, +delete_last_value(key)+ - removes last value
1272
1395
  # * +prev(key)+, +succ(key)+ - returns first value of adjacent key
1273
1396
  # * +prev(key, last: true)+, +succ(key, last: true)+ - returns last value
1274
1397
  #
@@ -1302,9 +1425,9 @@ end
1302
1425
  # tree.size # => 3 (total key-value pairs)
1303
1426
  # tree.get(1) # => "first one" (first value)
1304
1427
  # tree.get(1, last: true) # => "second one" (last value)
1305
- # tree.get_all(1) # => ["first one", "second one"] (all values)
1428
+ # tree.values(1) # => ["first one", "second one"] (all values)
1306
1429
  #
1307
- # tree.delete_one(1) # removes only "first one"
1430
+ # tree.delete_value(1) # removes only "first one"
1308
1431
  # tree.get(1) # => "second one"
1309
1432
  #
1310
1433
  # tree.delete(1) # removes all remaining values for key 1
@@ -1312,16 +1435,46 @@ end
1312
1435
  # @author Masahito Suzuki
1313
1436
  # @since 0.1.2
1314
1437
  class MultiRBTree < RBTree
1315
- def min(last: false)
1316
- return nil if @min_node == @nil_node || @min_node.value.empty?
1317
- [@min_node.key, last ? @min_node.value.last : @min_node.value.first]
1438
+ def initialize(*args, **kwargs)
1439
+ @value_count = 0
1440
+ super
1318
1441
  end
1319
1442
 
1320
- def max(last: false)
1321
- n = rightmost(@root)
1322
- return nil if n == @nil_node || n.value.empty?
1323
- [n.key, last ? n.value.last : n.value.first]
1324
- end
1443
+ # Returns the number of values stored in the tree.
1444
+ # @return [Integer] the number of values in the tree
1445
+ def size = @value_count
1446
+
1447
+ # Returns the minimum key-value pair without removing it.
1448
+ #
1449
+ # @param last [Boolean] whether to return the last value (default: false)
1450
+ # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
1451
+ # @example
1452
+ # tree = MultiRBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
1453
+ # tree.min # => [1, "one"]
1454
+ def min(last: false) = (n = min_node) && [n.key, n.value.send(last ? :last : :first)]
1455
+
1456
+ # Returns the maximum key-value pair without removing it.
1457
+ #
1458
+ # @param last [Boolean] whether to return the last value (default: false)
1459
+ # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
1460
+ # @example
1461
+ # tree = MultiRBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
1462
+ # tree.max # => [3, "three"]
1463
+ def max(last: false) = (n = max_node) && [n.key, n.value.send(last ? :last : :first)]
1464
+
1465
+ # Returns the last key-value pair without removing it.
1466
+ #
1467
+ # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
1468
+ # @example
1469
+ # tree = MultiRBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
1470
+ # tree.last # => [3, "three"]
1471
+ def last = max(last: true)
1472
+
1473
+ # Returns the number of values for a given key or the total number of key-value pairs if no key is given.
1474
+ #
1475
+ # @param key [Object, nil] the key to look up, or nil for total count
1476
+ # @return [Integer] the number of values for the key, or total count if no key is given
1477
+ def value_count(key = nil) = !key ? size : (@hash_index[key]&.value&.size || 0)
1325
1478
 
1326
1479
  # Retrieves a value associated with the given key.
1327
1480
  #
@@ -1334,41 +1487,66 @@ class MultiRBTree < RBTree
1334
1487
  # tree.insert(1, 'second')
1335
1488
  # tree.get(1) # => "first"
1336
1489
  # tree.get(1, last: true) # => "second"
1337
- def get(key, last: false)
1338
- n = find_node(key)
1339
- return nil if n == @nil_node || n.value.empty?
1340
- last ? n.value.last : n.value.first
1341
- end
1490
+ def value(key, last: false) = @hash_index[key]&.value&.send(last ? :last : :first)
1342
1491
 
1343
1492
  # Retrieves the first value associated with the given key.
1344
1493
  #
1345
1494
  # @param key [Object] the key to look up
1346
1495
  # @return [Object, nil] the first value for the key, or nil if not found
1347
- def get_first(key) = get(key, last: false)
1496
+ def first_value(key) = value(key)
1497
+ alias :get_first :first_value
1348
1498
 
1349
1499
  # Retrieves the last value associated with the given key.
1350
1500
  #
1351
1501
  # @param key [Object] the key to look up
1352
1502
  # @return [Object, nil] the last value for the key, or nil if not found
1353
- def get_last(key) = get(key, last: true)
1503
+ def last_value(key) = value(key, last: true)
1504
+ alias :get_last :last_value
1354
1505
 
1355
- def nearest(key, last: false)
1356
- n = find_nearest_node(key)
1357
- return nil if n == @nil_node || n.value.empty?
1358
- [n.key, last ? n.value.last : n.value.first]
1506
+ # Retrieves all values associated with the given key.
1507
+ #
1508
+ # @param key [Object] the key to look up
1509
+ # @return [Array, nil] an Array containing all values, or nil if not found
1510
+ # @example
1511
+ # tree = MultiRBTree.new
1512
+ # tree.insert(1, 'first')
1513
+ # tree.insert(1, 'second')
1514
+ # tree.values(1).to_a # => ["first", "second"]
1515
+ def values(key, reverse: false)
1516
+ return enum_for(:values, key) { value_count(key) } unless block_given?
1517
+ @hash_index[key]&.value&.send(reverse ? :reverse_each : :each) { |v| yield v }
1359
1518
  end
1519
+ alias :get_all :values
1360
1520
 
1361
- def prev(key, last: false)
1362
- pair = super(key)
1363
- return nil unless pair
1364
- [pair[0], last ? pair[1].last : pair[1].first]
1365
- end
1366
-
1367
- def succ(key, last: false)
1368
- pair = super(key)
1369
- return nil unless pair
1370
- [pair[0], last ? pair[1].last : pair[1].first]
1371
- end
1521
+ # Returns the nearest key-value pair without removing it.
1522
+ #
1523
+ # @param key [Object] the target key
1524
+ # @param last [Boolean] whether to return the last value (default: false)
1525
+ # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
1526
+ # @example
1527
+ # tree = MultiRBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
1528
+ # tree.nearest(4) # => [5, "five"]
1529
+ def nearest(key, last: false) = (pair = super(key)) && [pair[0], pair[1].send(last ? :last : :first)]
1530
+
1531
+ # Returns the previous key-value pair without removing it.
1532
+ #
1533
+ # @param key [Object] the target key
1534
+ # @param last [Boolean] whether to return the last value (default: false)
1535
+ # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
1536
+ # @example
1537
+ # tree = MultiRBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
1538
+ # tree.prev(4) # => [5, "five"]
1539
+ def prev(key, last: false) = (pair = super(key)) && [pair[0], pair[1].send(last ? :last : :first)]
1540
+
1541
+ # Returns the next key-value pair without removing it.
1542
+ #
1543
+ # @param key [Object] the target key
1544
+ # @param last [Boolean] whether to return the last value (default: false)
1545
+ # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
1546
+ # @example
1547
+ # tree = MultiRBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
1548
+ # tree.succ(4) # => [5, "five"]
1549
+ def succ(key, last: false) = (pair = super(key)) && [pair[0], pair[1].send(last ? :last : :first)]
1372
1550
 
1373
1551
  # Inserts a value for the given key.
1374
1552
  #
@@ -1377,15 +1555,16 @@ class MultiRBTree < RBTree
1377
1555
  #
1378
1556
  # @param key [Object] the key (must implement <=>)
1379
1557
  # @param value [Object] the value to insert
1558
+ # @param overwrite [Boolean] ignored for MultiRBTree which always appends
1380
1559
  # @return [Boolean] always returns true
1381
1560
  # @example
1382
1561
  # tree = MultiRBTree.new
1383
1562
  # tree.insert(1, 'first')
1384
1563
  # tree.insert(1, 'second') # adds another value for key 1
1385
- def insert(key, value)
1564
+ def insert_entry(key, value, **)
1386
1565
  if (node = @hash_index[key])
1387
1566
  node.value << value
1388
- @size += 1
1567
+ @value_count += 1
1389
1568
  return true
1390
1569
  end
1391
1570
  y = @nil_node
@@ -1395,7 +1574,7 @@ class MultiRBTree < RBTree
1395
1574
  cmp = key <=> x.key
1396
1575
  if cmp == 0
1397
1576
  x.value << value
1398
- @size += 1
1577
+ @value_count += 1
1399
1578
  return true
1400
1579
  elsif cmp < 0
1401
1580
  x = x.left
@@ -1416,7 +1595,7 @@ class MultiRBTree < RBTree
1416
1595
  z.right = @nil_node
1417
1596
  z.color = Node::RED
1418
1597
  insert_fixup(z)
1419
- @size += 1
1598
+ @value_count += 1
1420
1599
 
1421
1600
  if @min_node == @nil_node || (key <=> @min_node.key) < 0
1422
1601
  @min_node = z
@@ -1438,32 +1617,30 @@ class MultiRBTree < RBTree
1438
1617
  # tree = MultiRBTree.new
1439
1618
  # tree.insert(1, 'first')
1440
1619
  # tree.insert(1, 'second')
1441
- # tree.delete_one(1) # => "first"
1442
- # tree.delete_one(1, last: true) # => "second" (if more values existed)
1443
- def delete_one(key, last: false)
1444
- z = @hash_index[key] # O(1) lookup
1445
- return nil unless z
1446
-
1447
- value = last ? z.value.pop : z.value.shift
1448
- @size -= 1
1449
- if z.value.empty?
1450
- @hash_index.delete(key) # Remove from index when node removed
1451
- remove_node(z)
1452
- end
1620
+ # tree.delete_value(1) # => "first"
1621
+ # tree.delete_value(1, last: true) # => "second" (if more values existed)
1622
+ def delete_value(key, last: false)
1623
+ (z = @hash_index[key]) or return nil
1624
+ value = z.value.send(last ? :pop : :shift)
1625
+ z.value.empty? && delete_indexed_node(key)
1626
+ @value_count -= 1
1453
1627
  value
1454
1628
  end
1629
+ alias :delete_one :delete_value
1455
1630
 
1456
1631
  # Deletes the first value for the specified key.
1457
1632
  #
1458
1633
  # @param key [Object] the key to delete from
1459
1634
  # @return [Object, nil] the deleted value, or nil if key not found
1460
- def delete_first(key) = delete_one(key, last: false)
1635
+ def delete_first_value(key) = delete_value(key)
1636
+ alias :delete_first :delete_first_value
1461
1637
 
1462
1638
  # Deletes the last value for the specified key.
1463
1639
  #
1464
1640
  # @param key [Object] the key to delete from
1465
1641
  # @return [Object, nil] the deleted value, or nil if key not found
1466
- def delete_last(key) = delete_one(key, last: true)
1642
+ def delete_last_value(key) = delete_value(key, last: true)
1643
+ alias :delete_last :delete_last_value
1467
1644
 
1468
1645
  # Deletes all values for the specified key.
1469
1646
  #
@@ -1477,64 +1654,62 @@ class MultiRBTree < RBTree
1477
1654
  # tree.insert(1, 'second')
1478
1655
  # vals = tree.delete(1) # removes both values
1479
1656
  # vals.size # => 2
1480
- def delete(key)
1481
- z = @hash_index.delete(key) # O(1) lookup and remove from index
1482
- return nil unless z
1483
-
1484
- count = z.value.size
1485
- remove_node(z)
1486
- @size -= count
1487
- z.value
1657
+ def delete_key(key)
1658
+ return nil unless (z = @hash_index[key])
1659
+ @value_count -= (value = z.value).size
1660
+ delete_indexed_node(z.key)
1661
+ value
1488
1662
  end
1489
1663
 
1664
+ # Removes and returns the first key-value pair.
1665
+ #
1666
+ # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
1667
+ # @example
1668
+ # tree = MultiRBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
1669
+ # tree.shift # => [1, "one"]
1490
1670
  def shift
1491
- return nil if @min_node == @nil_node
1492
- node = @min_node
1493
- key = node.key
1494
- val = node.value.first
1495
- node.value.shift
1496
- @size -= 1
1497
- if node.value.empty?
1498
- delete_node(key)
1499
- end
1671
+ (key, vals = min_node&.pair) or return nil
1672
+ val = vals.shift
1673
+ vals.empty? && delete_indexed_node(key)
1674
+ @value_count -= 1
1500
1675
  [key, val]
1501
1676
  end
1502
1677
 
1503
- def pop
1504
- n = rightmost(@root)
1505
- return nil if n == @nil_node
1506
- key = n.key
1507
- val = n.value.last
1508
- n.value.pop
1509
- @size -= 1
1510
- if n.value.empty?
1511
- delete_node(key)
1512
- end
1513
- [key, val]
1514
- end
1515
-
1516
- # Retrieves all values associated with the given key.
1678
+ # Removes and returns the last key-value pair.
1517
1679
  #
1518
- # @param key [Object] the key to look up
1519
- # @return [Array, nil] an Array containing all values, or nil if not found
1680
+ # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
1520
1681
  # @example
1521
- # tree = MultiRBTree.new
1522
- # tree.insert(1, 'first')
1523
- # tree.insert(1, 'second')
1524
- # tree.get_all(1).to_a # => ["first", "second"]
1525
- def get_all(key)
1526
- return enum_for(:get_all, key) unless block_given?
1527
- @hash_index[key]&.value&.each { |v| yield v }
1682
+ # tree = MultiRBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
1683
+ # tree.pop # => [3, "three"]
1684
+ def pop
1685
+ (key, vals = max_node&.pair) or return nil
1686
+ val = vals.pop
1687
+ vals.empty? && delete_indexed_node(key)
1688
+ @value_count -= 1
1689
+ [key, val]
1528
1690
  end
1529
1691
 
1692
+ # @!visibility private
1530
1693
  private
1531
1694
 
1695
+ # Traverses the tree in ascending order, yielding each key-value pair.
1696
+ #
1697
+ # @param range [Range] the range of keys to traverse
1698
+ # @yield [Array(Object, Object)] each key-value pair
1699
+ # @yieldparam key [Object] the key
1700
+ # @yieldparam val [Object] the value
1532
1701
  def traverse_range_asc(...)
1533
- super { |k, vals| vals.each { |v| yield k, v } }
1702
+ super { |k, vals| vals.each { |v| yield [k, v] } }
1534
1703
  end
1535
1704
 
1705
+ # Traverses the tree in descending order, yielding each key-value pair.
1706
+ #
1707
+ # @param range [Range] the range of keys to traverse
1708
+ # @yield [Array(Object, Object)] each key-value pair
1709
+ # @yieldparam key [Object] the key
1710
+ # @yieldparam val [Object] the value
1536
1711
  def traverse_range_desc(...)
1537
- super { |k, vals| vals.reverse_each { |v| yield k, v } }
1712
+ super { |k, vals| vals.reverse_each { |v| yield [k, v] } }
1538
1713
  end
1539
1714
  end
1540
1715