The standard data structure to look for in this case is the disjoint-set data structure (also called union-find), https://en.wikipedia.org/wiki/Disjoint-set_data_structure. That should make your function fairly easy to implement.