13:- module(union_find_assoc, [
   14	union_find_assoc/2,
   15	make_set_assoc/3,
   16	union_assoc/4,
   17	union_all_assoc/3,
   18	find_assoc/4,
   19	find_assoc/5,
   20	disjoint_sets_assoc/2
   21]).   22
   23
   24
   25% union_find_assoc/2
   26% union_find_assoc(?UnionFind, +Elements)
   27%
   28% This predicate initializes a new ?UnionFind structure with a list of elements +Elements as keys.
   29union_find_assoc(UF, List) :-
   30	list_to_set(List, Set),
   31	empty_assoc(Assoc),
   32	union_find_assoc(1, Set, Assoc, UF).
   33
   34% union_find_assoc/3
   35% union_find_assoc(+LastID, +Set, +InitialAssoc, ?UnionFind)
   36%
   37% NOT EXPORTED
   38union_find_assoc(_, [], UF, UF).
   39union_find_assoc(I, [X|Xs], UF0, UF2) :-
   40	put_assoc(X, UF0, (X-0), UF1),
   41	succ(I, J),
   42	union_find_assoc(J, Xs, UF1, UF2).
   43
   44% make_set_assoc/3
   45% make_set_assoc(+UnionFindIn, +Element, ?UnionFindOut)
   46%
   47% This predicate makes a new set by creating a new element with a unique id +Element, a rank of 0, and a parent pointer
   48% to itself. The parent pointer to itself indicates that the element is the representative member of its own set.
   49make_set_assoc(UF0, X, UF1) :-
   50	\+get_assoc(X, UF0, _),
   51	put_assoc(X, UF0, (X-0), UF1).
   52
   53% union_assoc/4
   54% union_assoc(+UnionFindIn, +Element1, +Element2, ?UnionFindOut)
   55%
   56% This predicate uses find_assoc/5 to determine the roots of the trees +Element1 and +Element2 belong to.
   57% If the roots are distinct, the trees are combined by attaching the root of one to the root of the other.
   58% This predicate succeeds attaching the shorter tree (by rank) to the root of the taller tree in +UnionFindIn.
   59union_assoc(UF0, I, J, UF1) :-
   60	find_assoc(UF0, I, X, RankI, UF2),
   61	find_assoc(UF2, J, Y, RankJ, UF3),
   62	(X \== Y ->
   63		(RankI < RankJ -> put_assoc(X, UF3, Y-RankI, UF1) ; 
   64			(RankI > RankJ -> put_assoc(Y, UF3, X-RankJ, UF1) ; 
   65				put_assoc(Y, UF3, X-RankJ, UF4),
   66				succ(RankI, SrankI),
   67				put_assoc(X, UF4, X-SrankI, UF1))) ; UF1 = UF3).
   68
   69% union_all_assoc/3
   70% union_all_assoc(+UnionFindIn, +Elements, ?UnionFindOut)
   71%
   72% This predicate succeeds joining all the elements of the list +Elements in the union-find structure
   73% +UnionFindIn, producing the union-find structure ?UnionFindOut.
   74union_all_assoc(UF, [], UF).
   75union_all_assoc(UF, [_], UF).
   76union_all_assoc(UF0, [X,Y|Xs], UF2) :-
   77	union_assoc(UF0, X, Y, UF1),
   78	union_all_assoc(UF1, [Y|Xs], UF2).
   79
   80% find_assoc/4
   81% find_assoc(+UnionFindIn, +Element, ?Root, ?UnionFindOut)
   82%
   83% This predicate follows the chain of parent pointers from +Element up the tree until it reaches a ?Root element,
   84% whose parent is itself. ?Root is the representative member of the set to which ?Element belongs, and may be
   85% +Element itself. Path compression flattens the structure of the tree by making every node point to the root
   86% whenever find_assoc/4 is used on it.
   87find_assoc(UF0, I, X, UF1) :-
   88	get_assoc(I, UF0, J-R),
   89	(I == J -> X = J, UF1 = UF0 ; find_assoc(UF0, J, X, UF2), put_assoc(I, UF2, X-R, UF1)).
   90
   91% find_assoc/5
   92% find_assoc(+UnionFindIn, +Element, ?Root, ?Rank, ?UnionFindOut)
   93%
   94% Same as find_assoc/4, but returning also the ?Rank of the ?Root.
   95find_assoc(UF0, I, X, S, UF1) :-
   96	get_assoc(I, UF0, J-R),
   97	(I == J -> X = J, S = R, UF1 = UF0 ; find_assoc(UF0, J, X, S, UF2), put_assoc(I, UF2, X-R, UF1)).
   98
   99% disjoint_sets_assoc/2
  100% disjoint_sets_assoc(+UnionFind, ?Sets).
  101%
  102% This predicate succeeds when ?Sets is the list of disjoint sets on the +UnionFind structure.
  103disjoint_sets_assoc(UF, Sets) :-
  104	findall(Set, bagof(I, (Value,UF1)^(gen_assoc(I, UF, Value), find_assoc(UF, I, _, UF1)), Set), Sets)