From fab38945e391b5c92d66d4dcd9a73211728883f1 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 23 Mar 2026 12:31:14 +0530
Subject: [PATCH] feat: implement `Base.in` for `IntDisjointSet` and
 `DisjointSet`

---
 src/disjoint_set.jl       |  2 ++
 test/test_disjoint_set.jl | 25 +++++++++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/src/disjoint_set.jl b/src/disjoint_set.jl
index d4a396d7a..43db539ff 100644
--- a/src/disjoint_set.jl
+++ b/src/disjoint_set.jl
@@ -37,6 +37,7 @@ function Base.sizehint!(s::IntDisjointSet, n::Integer)
     sizehint!(s.ranks, n)
     return s
 end
+Base.in(i::Integer, s::IntDisjointSet) = checkbounds(Bool, s.parents, i)
 
 """
     num_groups(s::IntDisjointSet)
@@ -175,6 +176,7 @@ Base.iterate(s::DisjointSet) = iterate(s.revmap)
 Base.iterate(s::DisjointSet, i) = iterate(s.revmap, i)
 
 Base.length(s::DisjointSet) = length(s.internal)
+Base.in(el, s::DisjointSet) = haskey(s.intmap, el)
 
 """
     num_groups(s::DisjointSet)
diff --git a/test/test_disjoint_set.jl b/test/test_disjoint_set.jl
index a03465470..2b7c8e707 100644
--- a/test/test_disjoint_set.jl
+++ b/test/test_disjoint_set.jl
@@ -35,6 +35,17 @@
                     @test find_root!(s, T(3)) == T(2)
                 end
 
+                @testset "`Base.in`" begin
+                    @test 1 in s
+                    @test 1 in s2
+                    @test 10 in s
+                    @test 10 in s2
+                    @test !(11 in s)
+                    @test !(11 in s2)
+                    @test !(0 in s)
+                    @test !(0 in s2)
+                end
+
                 @testset "more tests" begin
                     # We cannot support arbitrary indexing and still use @inbounds with IntDisjointSet
                     # (and it's not useful anyway)
@@ -138,6 +149,20 @@
             @test find_root!(s, 7) == find_root!(s, 3)
         end
 
+        @testset "`Base.in`" begin
+            s = DisjointSet{Int}([1, 3, 5, 7, 9])
+            for i in [1, 3, 5, 7, 9]
+                @test i in s
+            end
+            @test !(2 in s)
+
+            s = DisjointSet{String}(["a", "b", "c", "e"])
+            for i in ["a", "b", "c", "e"]
+                @test i in s
+            end
+            @test !("d" in s)
+        end
+
         @testset "Some tests using non-integer disjoint sets" begin
             elems = ["a", "b", "c", "d"]
             a = DisjointSet{AbstractString}(elems)