feat: ✨ add an API to find the point with the max value in a user given dimension

Author: rishitc

comparable_utils.go

@@ -7,6 +7,13 @@ func min[T Comparable[T]](lhs, rhs *T, tcd int) *T {
 	return rhs
 }
 
+func max[T Comparable[T]](lhs, rhs *T, tcd int) *T {
+	if (*lhs).Order(*rhs, tcd) == Greater {
+		return lhs
+	}
+	return rhs
+}
+
 func distance[T Comparable[T]](src, dst *T) int {
 	return (*src).Dist(*dst)
 }

kdtree.go

@@ -82,6 +82,17 @@ func (t *KDTree[T]) FindMin(targetDimension int) (T, bool) {
 	return *res, true
 }
 
+func (t *KDTree[T]) FindMax(targetDimension int) (T, bool) {
+	if t.root == nil || targetDimension >= t.dimensions {
+		return t.zeroVal, false
+	}
+	res := findMax(t.dimensions, targetDimension, 0, t.root)
+	if res == nil {
+		return t.zeroVal, false
+	}
+	return *res, true
+}
+
 func (t *KDTree[T]) NearestNeighbor(value T) (T, bool) {
 	res := nearestNeighbor(t.dimensions, &value, nil, 0, t.root)
 	if res == nil {
@@ -426,3 +437,32 @@ func findMin[T Comparable[T]](d, tcd, cd int, r *kdNode[T]) *T {
 		return min(lMin, min(rMin, &r.value, tcd), tcd)
 	}
 }
+
+func findMax[T Comparable[T]](d, tcd, cd int, r *kdNode[T]) *T {
+	if r == nil {
+		return nil
+	}
+
+	var lMax *T
+	var rMax *T
+	ncd := (cd + 1) % d
+	rMax = findMax(d, tcd, ncd, r.right)
+	if tcd != cd {
+		lMax = findMax(d, tcd, ncd, r.left)
+	}
+	if lMax == nil && rMax == nil {
+		return &r.value
+	} else if lMax == nil {
+		if (*rMax).Order(r.value, tcd) == Greater {
+			return rMax
+		}
+		return &r.value
+	} else if rMax == nil {
+		if (*lMax).Order(r.value, tcd) == Greater {
+			return lMax
+		}
+		return &r.value
+	} else {
+		return max(lMax, max(rMax, &r.value, tcd), tcd)
+	}
+}

kdtree_2d_test.go

@@ -374,6 +374,72 @@ func Test2DFindMin2(t *testing.T) {
 	}
 }
 
+func Test2DFindMax1(t *testing.T) {
+	const dimensions = 2
+	ps := []Tensor2D{
+		{35, 90},
+		{60, 80},
+		{51, 75},
+		{70, 70},
+		{50, 50},
+		{25, 40},
+		{10, 30},
+		{1, 10},
+		{55, 1},
+	}
+	tree := kdtree.NewKDTreeWithValues(dimensions, ps)
+	testTable := []struct {
+		input    int
+		expected Tensor2D
+	}{
+		{
+			input:    0,
+			expected: Tensor2D{70, 70},
+		},
+		{
+			input:    1,
+			expected: Tensor2D{35, 90},
+		},
+	}
+	for _, v := range testTable {
+		nn, ok := tree.FindMax(v.input)
+		if !ok || !slices.Equal(nn[:], v.expected[:]) {
+			t.Fatalf("Expected closest point: %v, got %v", v.expected, nn)
+		}
+	}
+}
+
+func Test2DFindMax2(t *testing.T) {
+	const dimensions = 2
+	ps := []Tensor2D{
+		{35, 90},
+		{60, 80},
+		{51, 75},
+		{70, 70},
+		{50, 50},
+		{25, 40},
+		{10, 30},
+		{1, 10},
+		{55, 1},
+	}
+	tree := kdtree.NewKDTreeWithValues(dimensions, ps)
+	testTable := []struct {
+		input    int
+		expected Tensor2D
+	}{
+		{
+			input:    1,
+			expected: Tensor2D{55, 1},
+		},
+	}
+	for _, v := range testTable {
+		nn, ok := tree.FindMin(v.input)
+		if !ok || !slices.Equal(nn[:], v.expected[:]) {
+			t.Fatalf("Expected closest point: %v, got %v", v.expected, nn)
+		}
+	}
+}
+
 func Test2DDeleteAllNodesInTree(t *testing.T) {
 	treeNodes := kdtree.NewKDNode(Tensor2D{5, 6})
 	tree := kdtree.NewTestKDTree(2, treeNodes)