Union find の亜種として、Quick find と呼ばれているものがあり、これを使うと計算量に log がつく代わりに、特に今回の場合に実装がたいへん易しくなります。

（このあたりのお名前、信頼できる資料が見つからなく、自信がありませんため、ご存じの方は教えていただけると幸いです。）

Quick find について

Quick find は union find 同様、\([0, N[\) の集合分割を管理します。この集合分割の各要素を今後、グループと呼びます。

各グループに対して leader を定め、

• leader[i]: i の属するグループのリーダー
• members[i] リーダー i の属するグループ。リーダーでない場合は空

を管理し、union-by-size により更新します。すると union はクエリ数 \(Q\) として \(O(Q + N \log N)\) 時間で実行できます。

struct QuickFind {
    leader: Vec<usize>,
    members: Vec<Vec<usize>>,
}

差分管理について

各要素について、それとリーダーとの差分を管理すると良いです。

struct QuickFind {
    leader: Vec<usize>,
    members: Vec<Vec<usize>>,
    diff: Vec<i64>,
}