11// Copyright 2018-2025 the Deno authors. MIT license.
22import { assertEquals , assertStrictEquals , assertThrows } from "@std/assert" ;
3+ import { BinarySearchNode } from "./_binary_search_node.ts" ;
4+ import { RedBlackNode } from "./_red_black_node.ts" ;
35import { RedBlackTree } from "./red_black_tree.ts" ;
46import { ascend , descend } from "./comparators.ts" ;
57import { type Container , MyMath } from "./_test_utils.ts" ;
@@ -252,32 +254,40 @@ Deno.test("RedBlackTree works as exepcted with descend comparator", () => {
252254 }
253255} ) ;
254256
257+ function callback ( n : BinarySearchNode < Container > ) {
258+ let total_size = 1 ;
259+ total_size += n . left ?. value . st_size || 0 ;
260+ total_size += n . right ?. value . st_size || 0 ;
261+ n . value . st_size = total_size ;
262+ }
263+
255264Deno . test ( "RedBlackTree works with object items" , ( ) => {
256265 const tree : RedBlackTree < Container > = new RedBlackTree ( (
257266 a : Container ,
258267 b : Container ,
259- ) => ascend ( a . id , b . id ) ) ;
268+ ) => ascend ( a . id , b . id ) , callback ) ;
260269 const ids : number [ ] = [ - 10 , 9 , - 1 , 100 , 1 , 0 , - 100 , 10 , - 9 ] ;
261270
262271 for ( const [ i , id ] of ids . entries ( ) ) {
263- const newContainer : Container = { id, values : [ ] } ;
272+ const newContainer : Container = { id, values : [ ] , st_size : 1 } ;
264273 assertEquals ( tree . find ( newContainer ) , null ) ;
265274 assertEquals ( tree . insert ( newContainer ) , true ) ;
266275 newContainer . values . push ( i - 1 , i , i + 1 ) ;
267- assertStrictEquals ( tree . find ( { id, values : [ ] } ) , newContainer ) ;
276+ assertStrictEquals ( tree . find ( { id, values : [ ] , st_size : 1 } ) , newContainer ) ;
268277 assertEquals ( tree . size , i + 1 ) ;
278+ assertEquals ( tree . getRoot ( ) ?. value . st_size , i + 1 ) ;
269279 assertEquals ( tree . isEmpty ( ) , false ) ;
270280 }
271281 for ( const [ i , id ] of ids . entries ( ) ) {
272- const newContainer : Container = { id, values : [ ] } ;
282+ const newContainer : Container = { id, values : [ ] , st_size : 1 } ;
273283 assertEquals (
274- tree . find ( { id } as Container ) ,
275- { id , values : [ i - 1 , i , i + 1 ] } ,
284+ tree . find ( { id } as Container ) ?. id ,
285+ id ,
276286 ) ;
277287 assertEquals ( tree . insert ( newContainer ) , false ) ;
278288 assertEquals (
279- tree . find ( { id, values : [ ] } ) ,
280- { id , values : [ i - 1 , i , i + 1 ] } ,
289+ tree . find ( { id, values : [ ] , st_size : 1 } ) ?. id ,
290+ id ,
281291 ) ;
282292 assertEquals ( tree . size , ids . length ) ;
283293 assertEquals ( tree . isEmpty ( ) , false ) ;
@@ -295,18 +305,19 @@ Deno.test("RedBlackTree works with object items", () => {
295305 assertEquals ( tree . size , ids . length - i ) ;
296306 assertEquals ( tree . isEmpty ( ) , false ) ;
297307 assertEquals (
298- tree . find ( { id, values : [ ] } ) ,
299- { id , values : [ i - 1 , i , i + 1 ] } ,
308+ tree . find ( { id, values : [ ] , st_size : 1 } ) ?. id ,
309+ id ,
300310 ) ;
301311
302- assertEquals ( tree . remove ( { id, values : [ ] } ) , true ) ;
312+ assertEquals ( tree . remove ( { id, values : [ ] , st_size : 1 } ) , true ) ;
313+ assertEquals ( tree . getRoot ( ) ?. value . st_size || 0 , ids . length - i - 1 ) ;
303314 expected . splice ( expected . indexOf ( id ) , 1 ) ;
304315 assertEquals ( [ ...tree ] . map ( ( container ) => container . id ) , expected ) ;
305- assertEquals ( tree . find ( { id, values : [ ] } ) , null ) ;
316+ assertEquals ( tree . find ( { id, values : [ ] , st_size : 1 } ) , null ) ;
306317
307- assertEquals ( tree . remove ( { id, values : [ ] } ) , false ) ;
318+ assertEquals ( tree . remove ( { id, values : [ ] , st_size : 1 } ) , false ) ;
308319 assertEquals ( [ ...tree ] . map ( ( container ) => container . id ) , expected ) ;
309- assertEquals ( tree . find ( { id, values : [ ] } ) , null ) ;
320+ assertEquals ( tree . find ( { id, values : [ ] , st_size : 1 } ) , null ) ;
310321 }
311322 assertEquals ( tree . size , 0 ) ;
312323 assertEquals ( tree . isEmpty ( ) , true ) ;
0 commit comments