"CentroidCellID()" method for CellUnion

s2/cellunion.go

@@ -20,6 +20,7 @@ import (
 	"slices"
 	"sort"
 
+	"github.com/golang/geo/r3"
 	"github.com/golang/geo/s1"
 )
 
@@ -334,6 +335,33 @@ func (cu *CellUnion) CapBound() Cap {
 	return c
 }
 
+// CentroidCellID returns the CellID closest to the centroid of all cell centers.
+// Returns 0 (an invalid CellID) if the CellUnion is empty.
+func (cu CellUnion) CentroidCellID() CellID {
+	if len(cu) == 0 {
+		return 0
+	}
+
+	var sum r3.Vector
+	for _, cellID := range cu {
+		sum = sum.Add(CellFromCellID(cellID).Center().Vector)
+	}
+
+	target := Point{Vector: sum.Mul(1.0 / float64(len(cu))).Normalize()}
+
+	closest := cu[0]
+	distanceMin := CellFromCellID(closest).Center().Distance(target)
+	for _, cellID := range cu[1:] {
+		distance := CellFromCellID(cellID).Center().Distance(target)
+		if distance < distanceMin {
+			distanceMin = distance
+			closest = cellID
+		}
+	}
+
+	return closest
+}
+
 // ContainsCell reports whether this cell union contains the given cell.
 func (cu *CellUnion) ContainsCell(c Cell) bool {
 	return cu.ContainsCellID(c.id)

s2/cellunion_test.go

@@ -19,6 +19,7 @@ import (
 	"math"
 	"math/rand"
 	"reflect"
+	"slices"
 	"testing"
 
 	"github.com/golang/geo/r1"
@@ -656,6 +657,52 @@ func TestCellUnionRectBound(t *testing.T) {
 	}
 }
 
+func TestCellUnionCentroidCellID(t *testing.T) {
+	center := CellIDFromFace(0).ChildBeginAtLevel(10)
+	neighbors := center.AllNeighbors(11)
+
+	tests := []struct {
+		name string
+		cu   CellUnion
+		want CellID
+	}{
+		{
+			name: "empty",
+			cu:   CellUnion{},
+			want: 0,
+		},
+		{
+			name: "single cell",
+			cu:   CellUnion{center},
+			want: center,
+		},
+		{
+			name: "two adjacent cells",
+			cu:   CellUnion{center, neighbors[0]},
+		},
+		{
+			name: "three by three grid returns center",
+			cu:   append(neighbors, center),
+			want: center,
+		},
+	}
+
+	for _, test := range tests {
+		got := test.cu.CentroidCellID()
+		if test.want != 0 && got != test.want {
+			t.Errorf("%s: CentroidCellID() = %v, want %v", test.name, got, test.want)
+		}
+		if test.want == 0 && len(test.cu) > 0 {
+			// For non-empty unions without specific expected value,
+			// just verify the result is one of the input cells
+			found := slices.Contains(test.cu, got)
+			if !found {
+				t.Errorf("%s: CentroidCellID() = %v, not in input cells", test.name, got)
+			}
+		}
+	}
+}
+
 func TestCellUnionLeafCellsCovered(t *testing.T) {
 	fiveFaces := CellUnion{CellIDFromFace(0)}
 	fiveFaces.ExpandAtLevel(0)