ENH: Provide public API for fast DisjointSet subset size. #17714
Conversation
| result : int | ||
| Size of the subset containing `x`. | ||
| """ | ||
| return self._sizes[self[x]] |
There was a problem hiding this comment.
Other methods here are diligent about checking performing a
if x not in self._indices:
raise KeyError(x)before doing anything. Probably not a must, as the same KeyError will surface almost immediately anyway.
There was a problem hiding this comment.
Could be a good thing to have a better error message. But to follow our conventions it should be raise ValueError(message). And this would need a test with pytest.raises(..., match=message)
There was a problem hiding this comment.
I think the KeyError raised by __getitem__ is just as good. Note that merge and connected likewise do not perform explicit checking and rely on the same check in __getitem__.
There was a problem hiding this comment.
Which does not necessarily mean it's perfect 😉 I will let @fuglede make the call.
There was a problem hiding this comment.
Good point, sounds good to me.
| @@ -210,6 +210,25 @@ def subset(self, x): | |||
| nxt = self._nbrs[nxt] | |||
| return set(result) | |||
|
|
|||
| def subset_size(self, x): | |||
There was a problem hiding this comment.
Could also make sense to include this in the examples on how to use the class, i.e. something like
>>> disjoint_set.subset('a')
{'a', 3, 'b'}
>>> disjoint_set.subset_size('a')
3
This information is already stored internally as `disjoint_set._sizes`; this PR simply adds a public method to access it.
|
@fuglede is this good to go for you? |
Yes. Lint check is failing but for a file that hasn't changed here. |
Reference issue
Closes #17285.
What does this implement/fix?
Provide public API to quickly get the size of a DisjointSet subset.
This information is already stored internally as
disjoint_set._sizes; this PR simply adds a public method to access it.Additional information