Allow manually changing `UnionFind` rank

[jacob-greenfield]

Summary

I would like to manually set the rank for UnionFind elements.

Motivation

For my use case, I am trying to implement a puzzle solver that decides whether line segments are blank or solid by eventually merging them into one of two sets: index 0 or 1, respectively. Higher indices represent unsolved segments.

Since every merge can change the representative of one of the sets involved, I want to give priority to sets 0 and 1. That way, I can simply check forest.find_mut(x) < 2 to see if a segment is solved, or e.g. check forest.find_mut(x) == 1 to specifically check if it is solid. I think a very simple way to accomplish this would be giving those sets a high rank. Currently, there is no way to do this, so the only way is checking e.g. (forest.find_mut(x) == forest.find_mut(0)) || (forest.find_mut(x) == forest.find_mut(1)), which is slow and complicated.

Details

I think the best way to do this would be exposing the rank field on UnionFind, just something simple like:

pub fn rank(&self) -> &[u8] {
    &self.rank
}

pub fn rank_mut(&mut self) -> &mut [u8] {
    &mut self.rank
}

Then in my code, I can simply write forest.rank_mut()[0] = u8::MAX; forest.rank_mut()[1] = u8::MAX.

Changing the rank cannot cause any undefined behavior. The only correctness issue I can think of is integer overflow, which would cause a panic in debug mode or wrap in release mode. Besides, in my case, if I was really worried about a bug or invalid puzzle causing a merge between 0 and 1, I could use u8::MAX - 1.

The rank is only used to choose the new representative when merging (currently, purely for optimization), and since either choice would be correct, UnionFind still works correctly regardless of rank. I suppose the last thing to consider is whether or not you want to stabilize the rank field, but I can't think of any reason not to.

I am willing to implement this myself since it is a trivial change.

[bluss]

This is a bit far outside of the focus of the petgraph crate and I'd recommend writing your own union find implementation for this use case.