diff --git a/s2/cellunion.go b/s2/cellunion.go index 36993c9..a7704b3 100644 --- a/s2/cellunion.go +++ b/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) diff --git a/s2/cellunion_test.go b/s2/cellunion_test.go index e0b1e91..3a17fc6 100644 --- a/s2/cellunion_test.go +++ b/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)