Recently I've been trying this problem. I happen to come across a solution where we use DSU to optimize brute force but my implementation gives WA and I can't seem to figure out if I'm missing any proof of correctness or my implementation is incorrect. Any suggestions are appreciated.