We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
ENH: Expose DisjointSet._sizes
### Is your feature request related to a problem? Please describe.

Currently, `scipy.cluster.hierarchy.DisjointSet` internally keeps track of the size of each subset in the `_sizes` dict, making it available in O(1) (well, one needs to canonicalize keys first, in O(alpha(n)) ~ O(1)), but does not expose that info; the only way to access it publically is to reconstruct the actual subset (`.subset()`) and compute its size (`len()`), which is much slower (that's proportional to the size of the subset).

An example of algorithm that benefits from this info is "An efficient algorithm for exact evaluation of stochastic watersheds", Malmberg & Luengo 2014 (algorithm 2 relies on the disjoint-set, and equation 3 is the one where the size of the subset is needed).

### Describe the solution you'd like.

Add a `DisjointSet().size(x)` API (perhaps `size_of_subset` if one wants to be super clear), which is `return self._sizes[self[x]]` (where `self[x]` canonicalizes the key).

### Describe alternatives you've considered.

N/A

### Additional context (e.g. screenshots, GIFs)

N/A

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `scipy/_lib/_disjoint_set.py`
Content:
```
1 """
2 Disjoint set data structure
3 """
4 
5 
6 class DisjointSet:
7     """ Disjoint set data structure for incremental connectivity queries.
8 
9     .. versionadded:: 1.6.0
10 
11     Attributes
12     ----------
13     n_subsets : int
14         The number of subsets.
15 
16     Methods
17     -------
18     add
19     merge
20     connected
21     subset
22     subsets
23     __getitem__
24 
25     Notes
26     -----
27     This class implements the disjoint set [1]_, also known as the *union-find*
28     or *merge-find* data structure. The *find* operation (implemented in
29     `__getitem__`) implements the *path halving* variant. The *merge* method
30     implements the *merge by size* variant.
31 
32     References
33     ----------
34     .. [1] https://en.wikipedia.org/wiki/Disjoint-set_data_structure
35 
36     Examples
37     --------
38     >>> from scipy.cluster.hierarchy import DisjointSet
39 
40     Initialize a disjoint set:
41 
42     >>> disjoint_set = DisjointSet([1, 2, 3, 'a', 'b'])
43 
44     Merge some subsets:
45 
46     >>> disjoint_set.merge(1, 2)
47     True
48     >>> disjoint_set.merge(3, 'a')
49     True
50     >>> disjoint_set.merge('a', 'b')
51     True
52     >>> disjoint_set.merge('b', 'b')
53     False
54 
55     Find root elements:
56 
57     >>> disjoint_set[2]
58     1
59     >>> disjoint_set['b']
60     3
61 
62     Test connectivity:
63 
64     >>> disjoint_set.connected(1, 2)
65     True
66     >>> disjoint_set.connected(1, 'b')
67     False
68 
69     List elements in disjoint set:
70 
71     >>> list(disjoint_set)
72     [1, 2, 3, 'a', 'b']
73 
74     Get the subset containing 'a':
75 
76     >>> disjoint_set.subset('a')
77     {'a', 3, 'b'}
78 
79     Get all subsets in the disjoint set:
80 
81     >>> disjoint_set.subsets()
82     [{1, 2}, {'a', 3, 'b'}]
83     """
84     def __init__(self, elements=None):
85         self.n_subsets = 0
86         self._sizes = {}
87         self._parents = {}
88         # _nbrs is a circular linked list which links connected elements.
89         self._nbrs = {}
90         # _indices tracks the element insertion order in `__iter__`.
91         self._indices = {}
92         if elements is not None:
93             for x in elements:
94                 self.add(x)
95 
96     def __iter__(self):
97         """Returns an iterator of the elements in the disjoint set.
98 
99         Elements are ordered by insertion order.
100         """
101         return iter(self._indices)
102 
103     def __len__(self):
104         return len(self._indices)
105 
106     def __contains__(self, x):
107         return x in self._indices
108 
109     def __getitem__(self, x):
110         """Find the root element of `x`.
111 
112         Parameters
113         ----------
114         x : hashable object
115             Input element.
116 
117         Returns
118         -------
119         root : hashable object
120             Root element of `x`.
121         """
122         if x not in self._indices:
123             raise KeyError(x)
124 
125         # find by "path halving"
126         parents = self._parents
127         while self._indices[x] != self._indices[parents[x]]:
128             parents[x] = parents[parents[x]]
129             x = parents[x]
130         return x
131 
132     def add(self, x):
133         """Add element `x` to disjoint set
134         """
135         if x in self._indices:
136             return
137 
138         self._sizes[x] = 1
139         self._parents[x] = x
140         self._nbrs[x] = x
141         self._indices[x] = len(self._indices)
142         self.n_subsets += 1
143 
144     def merge(self, x, y):
145         """Merge the subsets of `x` and `y`.
146 
147         The smaller subset (the child) is merged into the larger subset (the
148         parent). If the subsets are of equal size, the root element which was
149         first inserted into the disjoint set is selected as the parent.
150 
151         Parameters
152         ----------
153         x, y : hashable object
154             Elements to merge.
155 
156         Returns
157         -------
158         merged : bool
159             True if `x` and `y` were in disjoint sets, False otherwise.
160         """
161         xr = self[x]
162         yr = self[y]
163         if self._indices[xr] == self._indices[yr]:
164             return False
165 
166         sizes = self._sizes
167         if (sizes[xr], self._indices[yr]) < (sizes[yr], self._indices[xr]):
168             xr, yr = yr, xr
169         self._parents[yr] = xr
170         self._sizes[xr] += self._sizes[yr]
171         self._nbrs[xr], self._nbrs[yr] = self._nbrs[yr], self._nbrs[xr]
172         self.n_subsets -= 1
173         return True
174 
175     def connected(self, x, y):
176         """Test whether `x` and `y` are in the same subset.
177 
178         Parameters
179         ----------
180         x, y : hashable object
181             Elements to test.
182 
183         Returns
184         -------
185         result : bool
186             True if `x` and `y` are in the same set, False otherwise.
187         """
188         return self._indices[self[x]] == self._indices[self[y]]
189 
190     def subset(self, x):
191         """Get the subset containing `x`.
192 
193         Parameters
194         ----------
195         x : hashable object
196             Input element.
197 
198         Returns
199         -------
200         result : set
201             Subset containing `x`.
202         """
203         if x not in self._indices:
204             raise KeyError(x)
205 
206         result = [x]
207         nxt = self._nbrs[x]
208         while self._indices[nxt] != self._indices[x]:
209             result.append(nxt)
210             nxt = self._nbrs[nxt]
211         return set(result)
212 
213     def subsets(self):
214         """Get all the subsets in the disjoint set.
215 
216         Returns
217         -------
218         result : list
219             Subsets in the disjoint set.
220         """
221         result = []
222         visited = set()
223         for x in self:
224             if x not in visited:
225                 xset = self.subset(x)
226                 visited.update(xset)
227                 result.append(xset)
228         return result
229
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/scipy/_lib/_disjoint_set.py b/scipy/_lib/_disjoint_set.py
--- a/scipy/_lib/_disjoint_set.py
+++ b/scipy/_lib/_disjoint_set.py
@@ -19,6 +19,7 @@
     merge
     connected
     subset
+    subset_size
     subsets
     __getitem__
 
@@ -76,6 +77,12 @@
     >>> disjoint_set.subset('a')
     {'a', 3, 'b'}
 
+    Get the size of the subset containing 'a' (without actually instantiating
+    the subset):
+
+    >>> disjoint_set.subset_size('a')
+    3
+
     Get all subsets in the disjoint set:
 
     >>> disjoint_set.subsets()
@@ -210,6 +217,25 @@
             nxt = self._nbrs[nxt]
         return set(result)
 
+    def subset_size(self, x):
+        """Get the size of the subset containing `x`.
+
+        Note that this method is faster than ``len(self.subset(x))`` because
+        the size is directly read off an internal field, without the need to
+        instantiate the full subset.
+
+        Parameters
+        ----------
+        x : hashable object
+            Input element.
+
+        Returns
+        -------
+        result : int
+            Size of the subset containing `x`.
+        """
+        return self._sizes[self[x]]
+
     def subsets(self):
         """Get all the subsets in the disjoint set.

{"golden_diff": "diff --git a/scipy/_lib/_disjoint_set.py b/scipy/_lib/_disjoint_set.py\n--- a/scipy/_lib/_disjoint_set.py\n+++ b/scipy/_lib/_disjoint_set.py\n@@ -19,6 +19,7 @@\n     merge\n     connected\n     subset\n+    subset_size\n     subsets\n     __getitem__\n \n@@ -76,6 +77,12 @@\n     >>> disjoint_set.subset('a')\n     {'a', 3, 'b'}\n \n+    Get the size of the subset containing 'a' (without actually instantiating\n+    the subset):\n+\n+    >>> disjoint_set.subset_size('a')\n+    3\n+\n     Get all subsets in the disjoint set:\n \n     >>> disjoint_set.subsets()\n@@ -210,6 +217,25 @@\n             nxt = self._nbrs[nxt]\n         return set(result)\n \n+    def subset_size(self, x):\n+        \"\"\"Get the size of the subset containing `x`.\n+\n+        Note that this method is faster than ``len(self.subset(x))`` because\n+        the size is directly read off an internal field, without the need to\n+        instantiate the full subset.\n+\n+        Parameters\n+        ----------\n+        x : hashable object\n+            Input element.\n+\n+        Returns\n+        -------\n+        result : int\n+            Size of the subset containing `x`.\n+        \"\"\"\n+        return self._sizes[self[x]]\n+\n     def subsets(self):\n         \"\"\"Get all the subsets in the disjoint set.\n", "issue": "ENH: Expose DisjointSet._sizes\n### Is your feature request related to a problem? Please describe.\n\nCurrently, `scipy.cluster.hierarchy.DisjointSet` internally keeps track of the size of each subset in the `_sizes` dict, making it available in O(1) (well, one needs to canonicalize keys first, in O(alpha(n)) ~ O(1)), but does not expose that info; the only way to access it publically is to reconstruct the actual subset (`.subset()`) and compute its size (`len()`), which is much slower (that's proportional to the size of the subset).\r\n\r\nAn example of algorithm that benefits from this info is \"An efficient algorithm for exact evaluation of stochastic watersheds\", Malmberg & Luengo 2014 (algorithm 2 relies on the disjoint-set, and equation 3 is the one where the size of the subset is needed).\n\n### Describe the solution you'd like.\n\nAdd a `DisjointSet().size(x)` API (perhaps `size_of_subset` if one wants to be super clear), which is `return self._sizes[self[x]]` (where `self[x]` canonicalizes the key).\n\n### Describe alternatives you've considered.\n\nN/A\n\n### Additional context (e.g. screenshots, GIFs)\n\nN/A\n", "before_files": [{"content": "\"\"\"\nDisjoint set data structure\n\"\"\"\n\n\nclass DisjointSet:\n    \"\"\" Disjoint set data structure for incremental connectivity queries.\n\n    .. versionadded:: 1.6.0\n\n    Attributes\n    ----------\n    n_subsets : int\n        The number of subsets.\n\n    Methods\n    -------\n    add\n    merge\n    connected\n    subset\n    subsets\n    __getitem__\n\n    Notes\n    -----\n    This class implements the disjoint set [1]_, also known as the *union-find*\n    or *merge-find* data structure. The *find* operation (implemented in\n    `__getitem__`) implements the *path halving* variant. The *merge* method\n    implements the *merge by size* variant.\n\n    References\n    ----------\n    .. [1] https://en.wikipedia.org/wiki/Disjoint-set_data_structure\n\n    Examples\n    --------\n    >>> from scipy.cluster.hierarchy import DisjointSet\n\n    Initialize a disjoint set:\n\n    >>> disjoint_set = DisjointSet([1, 2, 3, 'a', 'b'])\n\n    Merge some subsets:\n\n    >>> disjoint_set.merge(1, 2)\n    True\n    >>> disjoint_set.merge(3, 'a')\n    True\n    >>> disjoint_set.merge('a', 'b')\n    True\n    >>> disjoint_set.merge('b', 'b')\n    False\n\n    Find root elements:\n\n    >>> disjoint_set[2]\n    1\n    >>> disjoint_set['b']\n    3\n\n    Test connectivity:\n\n    >>> disjoint_set.connected(1, 2)\n    True\n    >>> disjoint_set.connected(1, 'b')\n    False\n\n    List elements in disjoint set:\n\n    >>> list(disjoint_set)\n    [1, 2, 3, 'a', 'b']\n\n    Get the subset containing 'a':\n\n    >>> disjoint_set.subset('a')\n    {'a', 3, 'b'}\n\n    Get all subsets in the disjoint set:\n\n    >>> disjoint_set.subsets()\n    [{1, 2}, {'a', 3, 'b'}]\n    \"\"\"\n    def __init__(self, elements=None):\n        self.n_subsets = 0\n        self._sizes = {}\n        self._parents = {}\n        # _nbrs is a circular linked list which links connected elements.\n        self._nbrs = {}\n        # _indices tracks the element insertion order in `__iter__`.\n        self._indices = {}\n        if elements is not None:\n            for x in elements:\n                self.add(x)\n\n    def __iter__(self):\n        \"\"\"Returns an iterator of the elements in the disjoint set.\n\n        Elements are ordered by insertion order.\n        \"\"\"\n        return iter(self._indices)\n\n    def __len__(self):\n        return len(self._indices)\n\n    def __contains__(self, x):\n        return x in self._indices\n\n    def __getitem__(self, x):\n        \"\"\"Find the root element of `x`.\n\n        Parameters\n        ----------\n        x : hashable object\n            Input element.\n\n        Returns\n        -------\n        root : hashable object\n            Root element of `x`.\n        \"\"\"\n        if x not in self._indices:\n            raise KeyError(x)\n\n        # find by \"path halving\"\n        parents = self._parents\n        while self._indices[x] != self._indices[parents[x]]:\n            parents[x] = parents[parents[x]]\n            x = parents[x]\n        return x\n\n    def add(self, x):\n        \"\"\"Add element `x` to disjoint set\n        \"\"\"\n        if x in self._indices:\n            return\n\n        self._sizes[x] = 1\n        self._parents[x] = x\n        self._nbrs[x] = x\n        self._indices[x] = len(self._indices)\n        self.n_subsets += 1\n\n    def merge(self, x, y):\n        \"\"\"Merge the subsets of `x` and `y`.\n\n        The smaller subset (the child) is merged into the larger subset (the\n        parent). If the subsets are of equal size, the root element which was\n        first inserted into the disjoint set is selected as the parent.\n\n        Parameters\n        ----------\n        x, y : hashable object\n            Elements to merge.\n\n        Returns\n        -------\n        merged : bool\n            True if `x` and `y` were in disjoint sets, False otherwise.\n        \"\"\"\n        xr = self[x]\n        yr = self[y]\n        if self._indices[xr] == self._indices[yr]:\n            return False\n\n        sizes = self._sizes\n        if (sizes[xr], self._indices[yr]) < (sizes[yr], self._indices[xr]):\n            xr, yr = yr, xr\n        self._parents[yr] = xr\n        self._sizes[xr] += self._sizes[yr]\n        self._nbrs[xr], self._nbrs[yr] = self._nbrs[yr], self._nbrs[xr]\n        self.n_subsets -= 1\n        return True\n\n    def connected(self, x, y):\n        \"\"\"Test whether `x` and `y` are in the same subset.\n\n        Parameters\n        ----------\n        x, y : hashable object\n            Elements to test.\n\n        Returns\n        -------\n        result : bool\n            True if `x` and `y` are in the same set, False otherwise.\n        \"\"\"\n        return self._indices[self[x]] == self._indices[self[y]]\n\n    def subset(self, x):\n        \"\"\"Get the subset containing `x`.\n\n        Parameters\n        ----------\n        x : hashable object\n            Input element.\n\n        Returns\n        -------\n        result : set\n            Subset containing `x`.\n        \"\"\"\n        if x not in self._indices:\n            raise KeyError(x)\n\n        result = [x]\n        nxt = self._nbrs[x]\n        while self._indices[nxt] != self._indices[x]:\n            result.append(nxt)\n            nxt = self._nbrs[nxt]\n        return set(result)\n\n    def subsets(self):\n        \"\"\"Get all the subsets in the disjoint set.\n\n        Returns\n        -------\n        result : list\n            Subsets in the disjoint set.\n        \"\"\"\n        result = []\n        visited = set()\n        for x in self:\n            if x not in visited:\n                xset = self.subset(x)\n                visited.update(xset)\n                result.append(xset)\n        return result\n", "path": "scipy/_lib/_disjoint_set.py"}], "after_files": [{"content": "\"\"\"\nDisjoint set data structure\n\"\"\"\n\n\nclass DisjointSet:\n    \"\"\" Disjoint set data structure for incremental connectivity queries.\n\n    .. versionadded:: 1.6.0\n\n    Attributes\n    ----------\n    n_subsets : int\n        The number of subsets.\n\n    Methods\n    -------\n    add\n    merge\n    connected\n    subset\n    subset_size\n    subsets\n    __getitem__\n\n    Notes\n    -----\n    This class implements the disjoint set [1]_, also known as the *union-find*\n    or *merge-find* data structure. The *find* operation (implemented in\n    `__getitem__`) implements the *path halving* variant. The *merge* method\n    implements the *merge by size* variant.\n\n    References\n    ----------\n    .. [1] https://en.wikipedia.org/wiki/Disjoint-set_data_structure\n\n    Examples\n    --------\n    >>> from scipy.cluster.hierarchy import DisjointSet\n\n    Initialize a disjoint set:\n\n    >>> disjoint_set = DisjointSet([1, 2, 3, 'a', 'b'])\n\n    Merge some subsets:\n\n    >>> disjoint_set.merge(1, 2)\n    True\n    >>> disjoint_set.merge(3, 'a')\n    True\n    >>> disjoint_set.merge('a', 'b')\n    True\n    >>> disjoint_set.merge('b', 'b')\n    False\n\n    Find root elements:\n\n    >>> disjoint_set[2]\n    1\n    >>> disjoint_set['b']\n    3\n\n    Test connectivity:\n\n    >>> disjoint_set.connected(1, 2)\n    True\n    >>> disjoint_set.connected(1, 'b')\n    False\n\n    List elements in disjoint set:\n\n    >>> list(disjoint_set)\n    [1, 2, 3, 'a', 'b']\n\n    Get the subset containing 'a':\n\n    >>> disjoint_set.subset('a')\n    {'a', 3, 'b'}\n\n    Get the size of the subset containing 'a' (without actually instantiating\n    the subset):\n\n    >>> disjoint_set.subset_size('a')\n    3\n\n    Get all subsets in the disjoint set:\n\n    >>> disjoint_set.subsets()\n    [{1, 2}, {'a', 3, 'b'}]\n    \"\"\"\n    def __init__(self, elements=None):\n        self.n_subsets = 0\n        self._sizes = {}\n        self._parents = {}\n        # _nbrs is a circular linked list which links connected elements.\n        self._nbrs = {}\n        # _indices tracks the element insertion order in `__iter__`.\n        self._indices = {}\n        if elements is not None:\n            for x in elements:\n                self.add(x)\n\n    def __iter__(self):\n        \"\"\"Returns an iterator of the elements in the disjoint set.\n\n        Elements are ordered by insertion order.\n        \"\"\"\n        return iter(self._indices)\n\n    def __len__(self):\n        return len(self._indices)\n\n    def __contains__(self, x):\n        return x in self._indices\n\n    def __getitem__(self, x):\n        \"\"\"Find the root element of `x`.\n\n        Parameters\n        ----------\n        x : hashable object\n            Input element.\n\n        Returns\n        -------\n        root : hashable object\n            Root element of `x`.\n        \"\"\"\n        if x not in self._indices:\n            raise KeyError(x)\n\n        # find by \"path halving\"\n        parents = self._parents\n        while self._indices[x] != self._indices[parents[x]]:\n            parents[x] = parents[parents[x]]\n            x = parents[x]\n        return x\n\n    def add(self, x):\n        \"\"\"Add element `x` to disjoint set\n        \"\"\"\n        if x in self._indices:\n            return\n\n        self._sizes[x] = 1\n        self._parents[x] = x\n        self._nbrs[x] = x\n        self._indices[x] = len(self._indices)\n        self.n_subsets += 1\n\n    def merge(self, x, y):\n        \"\"\"Merge the subsets of `x` and `y`.\n\n        The smaller subset (the child) is merged into the larger subset (the\n        parent). If the subsets are of equal size, the root element which was\n        first inserted into the disjoint set is selected as the parent.\n\n        Parameters\n        ----------\n        x, y : hashable object\n            Elements to merge.\n\n        Returns\n        -------\n        merged : bool\n            True if `x` and `y` were in disjoint sets, False otherwise.\n        \"\"\"\n        xr = self[x]\n        yr = self[y]\n        if self._indices[xr] == self._indices[yr]:\n            return False\n\n        sizes = self._sizes\n        if (sizes[xr], self._indices[yr]) < (sizes[yr], self._indices[xr]):\n            xr, yr = yr, xr\n        self._parents[yr] = xr\n        self._sizes[xr] += self._sizes[yr]\n        self._nbrs[xr], self._nbrs[yr] = self._nbrs[yr], self._nbrs[xr]\n        self.n_subsets -= 1\n        return True\n\n    def connected(self, x, y):\n        \"\"\"Test whether `x` and `y` are in the same subset.\n\n        Parameters\n        ----------\n        x, y : hashable object\n            Elements to test.\n\n        Returns\n        -------\n        result : bool\n            True if `x` and `y` are in the same set, False otherwise.\n        \"\"\"\n        return self._indices[self[x]] == self._indices[self[y]]\n\n    def subset(self, x):\n        \"\"\"Get the subset containing `x`.\n\n        Parameters\n        ----------\n        x : hashable object\n            Input element.\n\n        Returns\n        -------\n        result : set\n            Subset containing `x`.\n        \"\"\"\n        if x not in self._indices:\n            raise KeyError(x)\n\n        result = [x]\n        nxt = self._nbrs[x]\n        while self._indices[nxt] != self._indices[x]:\n            result.append(nxt)\n            nxt = self._nbrs[nxt]\n        return set(result)\n\n    def subset_size(self, x):\n        \"\"\"Get the size of the subset containing `x`.\n\n        Note that this method is faster than ``len(self.subset(x))`` because\n        the size is directly read off an internal field, without the need to\n        instantiate the full subset.\n\n        Parameters\n        ----------\n        x : hashable object\n            Input element.\n\n        Returns\n        -------\n        result : int\n            Size of the subset containing `x`.\n        \"\"\"\n        return self._sizes[self[x]]\n\n    def subsets(self):\n        \"\"\"Get all the subsets in the disjoint set.\n\n        Returns\n        -------\n        result : list\n            Subsets in the disjoint set.\n        \"\"\"\n        result = []\n        visited = set()\n        for x in self:\n            if x not in visited:\n                xset = self.subset(x)\n                visited.update(xset)\n                result.append(xset)\n        return result\n", "path": "scipy/_lib/_disjoint_set.py"}]}