diff --git a/README.md b/README.md index 0947d3f..841faa2 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,7 @@ type DSU[T comparable] interface { Find(x T) T // Find returns the representative element (root) of x Union(x, y T) bool // Union merges the sets containing x and y Connected(x, y T) bool // Connected reports whether x and y are in the same set + IsMember(x T) bool // IsMember reports whether x is present (no mutation) Groups() map[T][]T // Groups returns all connected components } ``` diff --git a/compact/compact.go b/compact/compact.go index 89ca6f0..cf6de16 100644 --- a/compact/compact.go +++ b/compact/compact.go @@ -100,6 +100,11 @@ func (dsu *DSU) Connected(x, y int) bool { return dsu.Find(x) == dsu.Find(y) } +// IsMember reports whether x is within the managed range [0, n). +func (dsu *DSU) IsMember(x int) bool { + return dsu.boundsCheck(x) +} + // Groups returns a map from root -> slice of elements in that set. func (dsu *DSU) Groups() map[int][]int { groups := make(map[int][]int) diff --git a/compact/compact_test.go b/compact/compact_test.go index 6a974b5..b85d260 100644 --- a/compact/compact_test.go +++ b/compact/compact_test.go @@ -40,6 +40,16 @@ func TestCompactGroups(t *testing.T) { } } +func TestCompactIsMember(t *testing.T) { + dsu := New(3) + if !dsu.IsMember(0) || !dsu.IsMember(2) { + t.Fatalf("expected in-range elements to be members") + } + if dsu.IsMember(-1) || dsu.IsMember(3) { + t.Fatalf("expected out-of-range elements to report not a member") + } +} + // TestNewNegativeSize ensures New() dos not panic when called with a negative size. func TestCompactNewNegativeSize(t *testing.T) { _ = New(-5) diff --git a/gdsu.go b/gdsu.go index b9dcec4..8c587f4 100644 --- a/gdsu.go +++ b/gdsu.go @@ -12,6 +12,7 @@ package gdsu // - Find: return the canonical representative (root) of x's set. // - Union: merge sets containing x and y, return true if they were separate. // - Connected: report whether x and y belong to the same set. +// - IsMember: report whether x has been registered with the DSU. // - Groups: return all current sets as root -> slice of elements. type DSU[T comparable] interface { // Find returns the representative (root) of the set containing x. @@ -23,6 +24,9 @@ type DSU[T comparable] interface { // Connected returns true iff x and y belong to the same set. Connected(x, y T) bool + // IsMember reports whether x is present in the DSU without mutating the DSU. + IsMember(x T) bool + // Groups returns a mapping of each root to all elements in its set. Groups() map[T][]T } diff --git a/gdsu_test.go b/gdsu_test.go index e83ad4f..dc104df 100644 --- a/gdsu_test.go +++ b/gdsu_test.go @@ -9,6 +9,7 @@ type fakeDSU[T comparable] struct{} func (f *fakeDSU[T]) Find(x T) T { return x } func (f *fakeDSU[T]) Union(x, y T) bool { return true } func (f *fakeDSU[T]) Connected(x, y T) bool { return true } +func (f *fakeDSU[T]) IsMember(x T) bool { return true } func (f *fakeDSU[T]) Groups() map[T][]T { return map[T][]T{} } func TestInterfaceCompiles(t *testing.T) { diff --git a/sparse/sparse.go b/sparse/sparse.go index 0537577..c413eac 100644 --- a/sparse/sparse.go +++ b/sparse/sparse.go @@ -87,6 +87,12 @@ func (dsu *DSU[T]) Connected(x, y T) bool { return dsu.Find(x) == dsu.Find(y) } +// IsMember reports whether x has been seen by the DSU without adding it. +func (dsu *DSU[T]) IsMember(x T) bool { + _, exists := dsu.parent[x] + return exists +} + // Groups returns a map from root -> slice of elements in that set. func (dsu *DSU[T]) Groups() map[T][]T { groups := make(map[T][]T) diff --git a/sparse/sparse_test.go b/sparse/sparse_test.go index 5773773..ba3d704 100644 --- a/sparse/sparse_test.go +++ b/sparse/sparse_test.go @@ -66,6 +66,22 @@ func TestSparseFindCreatesNew(t *testing.T) { } } +func TestSparseIsMember(t *testing.T) { + dsu := New[int](1, 2) + + if !dsu.IsMember(1) || !dsu.IsMember(2) { + t.Fatalf("expected initial elements to be members") + } + if dsu.IsMember(3) { + t.Fatalf("expected unseen element to report not a member") + } + // Calling IsMember on an unseen element must not mutate the DSU. + groups := dsu.Groups() + if len(groups) != 2 { + t.Fatalf("expected 2 singleton groups, got %d", len(groups)) + } +} + // TestSparseUnionOnNewElements ensures Union() works even if elements were never added before. func TestSparseUnionOnNewElements(t *testing.T) { dsu := New[int]()