@@ -402,13 +402,35 @@ func (g *Graph[K]) Add(nodes ...Node[K]) {
402402
403403 // Invariant check: the node should have been added to the graph.
404404 if g .Len () != preLen + 1 {
405- panic ("node not added" )
405+ if len (g .layers ) > 0 && g .layers [len (g .layers )- 1 ].entry () == nil {
406+ g .layers = g .layers [:len (g .layers )- 1 ]
407+ }
406408 }
407409 }
408410}
409411
410412// Search finds the k nearest neighbors from the target node.
411413func (h * Graph [K ]) Search (near Vector , k int ) []Node [K ] {
414+ sr := h .search (near , k )
415+ out := make ([]Node [K ], len (sr ))
416+ for i , node := range sr {
417+ out [i ] = node .Node
418+ }
419+ return out
420+ }
421+
422+ // SearchWithDistance finds the k nearest neighbors from the target node
423+ // and returns the distance.
424+ func (h * Graph [K ]) SearchWithDistance (near Vector , k int ) []SearchResult [K ] {
425+ return h .search (near , k )
426+ }
427+
428+ type SearchResult [T cmp.Ordered ] struct {
429+ Node [T ]
430+ Distance float32
431+ }
432+
433+ func (h * Graph [K ]) search (near Vector , k int ) []SearchResult [K ] {
412434 h .assertDims (near )
413435 if len (h .layers ) == 0 {
414436 return nil
@@ -434,10 +456,13 @@ func (h *Graph[K]) Search(near Vector, k int) []Node[K] {
434456 }
435457
436458 nodes := searchPoint .search (k , efSearch , near , h .Distance )
437- out := make ([]Node [K ], 0 , len (nodes ))
459+ out := make ([]SearchResult [K ], 0 , len (nodes ))
438460
439461 for _ , node := range nodes {
440- out = append (out , node .node .Node )
462+ out = append (out , SearchResult [K ]{
463+ Node : node .node .Node ,
464+ Distance : node .dist ,
465+ })
441466 }
442467
443468 return out
@@ -462,17 +487,33 @@ func (h *Graph[K]) Delete(key K) bool {
462487 return false
463488 }
464489
490+ var deleteLayer = map [int ]struct {}{}
465491 var deleted bool
466- for _ , layer := range h .layers {
492+ for i , layer := range h .layers {
467493 node , ok := layer .nodes [key ]
468494 if ! ok {
469495 continue
470496 }
471497 delete (layer .nodes , key )
498+ if len (layer .nodes ) == 0 {
499+ deleteLayer [i ] = struct {}{}
500+ }
472501 node .isolate (h .M )
473502 deleted = true
474503 }
475504
505+ if len (deleteLayer ) > 0 {
506+ var newLayers = make ([]* layer [K ], 0 , len (h .layers )- len (deleteLayer ))
507+ for i , layer := range h .layers {
508+ if _ , ok := deleteLayer [i ]; ok {
509+ continue
510+ }
511+ newLayers = append (newLayers , layer )
512+ }
513+
514+ h .layers = newLayers
515+ }
516+
476517 return deleted
477518}
478519
0 commit comments