For a final project I used a weighted union-find and decided to implement it in Python. I wrote it to be as general as possible so it's very portable. I do plan on looking into path compression to make it better, but that's in the future. Otherwise, I have plans to make it more generic so that it's even more portable. Before I do so, I'd just like to get some more eyes on it to see if there are optimizations I should make or if there are security flaws.

class WUF:
    """This class implements the Weighted Union Find (WUF) data structure for contact tracing."""

    Pair = list[str, int]
    """A data type used to combine the root and size information into a single location."""

    def __init__(self, iTree: dict[str, Pair] = None):
        """initializes the WUF.

        A fairly standard __init__ function to initliaze the data necessary for the WUF.
        If a tree is not provided, an empty tree is created.

        :param self: self
        :param iTree: the tree structure to construct the WUF from
        :type iTree: dict[str, (str, int)]
        if not iTree:
            iTree = {}

        self.tree: dict[str, WUF.Pair] = iTree
        """Used to store the WUF tree with the student ID and it's corresponding root information."""

    def __eq__(self, other: "WUF") -> bool:
        """returns if two WUF objects have the same tree

        :param self: self
        :param other: the WUF object to compare to
        :type other: WUF
        :returns: if the two trees are equal
        :rtype: bool
        return self.tree == other.tree

    def __len__(self) -> int:
        """returns the size of the tree

        :param self: self
        :returns: the length of the tree
        :rtype: int

        return len(self.tree)

    def __iter__(self):
        """returns the tree to use as an iterable

        :param self: self
        :returns: the tree iterator
        return self.tree.__iter__()

    def __getitem__(self, item: int) -> Pair:
        """returns the value in tree at index item

        :param self: self
        :param item: the index to find the value at
        :type item: int
        :returns: the value in tree at item
        :rtype: list[str, int]
        return self.tree[item]

    def addRoot(self, uid: str) -> bool:
        """returns whether the operation was successful

        Adds a new root to the tree with its root set to self (meaning it's top level) and with a size of one.

        :param self: self
        :param uid: the Student ID to add to the tree
        :type uid: str
        :returns: whether the operation succeeded
        :rtype: bool
        if uid not in self.tree:
            self.tree[uid] = [uid, 1]
            return True

        return False

    def getRoot(self, uid: str) -> str:
        """returns the root of the union the id is in

        This function traverses the tree in order to find the root of the union that ID is a part of.

        :param self: self
        :param uid: the Student ID to retrieve the union for
        :type uid: str
        :returns: the root student ID in the union
        :rtype: str
        if uid not in self.tree:
            raise IndexError(f"{uid} is not in the tree")

        root = uid
        while root != self.tree[root][0]:
            root = self.tree[root][0]

        return root

    def getRootGroups(self) -> dict[str, list[str]]:
        """returns the tree in the format of root groups

        This function will parse the tree and return it as root groups such that the top level roots are the dictionary
        keys and a list of UIDs is the value.

        :param self: self
        :returns: the root group version of the tree
        :rtype: dict[str, list[str]]
        rootGroups: dict[str, list[str]] = {}

        for key in self.tree:
            root: str = self.getRoot(key)

            if root not in rootGroups:
                rootGroups[root] = []

            if key not in rootGroups[root]:

        return rootGroups

    def union(self, p: str, q: str) -> None:
        """unions two items in the tree

        Unions data in the tree using the WUF algorithm.
        If the size of q's root union is greater than the size of p's root union, then add p to q.
        Otherwise, q will be added to p.

        :param self: Self
        :param p: the first item to union
        :param q: the second item to union
        :type p: str
        :type q: str

        pRoot = self.getRoot(p)
        qRoot = self.getRoot(q)

        if pRoot == qRoot:

        if self.tree[qRoot][1] > self.tree[pRoot][1]:
            self.tree[pRoot][0] = qRoot
            if self.tree[pRoot][1] == self.tree[qRoot][1]:
                self.tree[pRoot][1] += self.tree[qRoot][1]

            self.tree[qRoot][0] = pRoot

1 Answer 1



Docstrings are for users of the class. Docstrings should explain what is being done and returned. Currently your docstrings explain how it is being done, which should be comments.

Currently as a user, I would have no idea what this data structure is, or how to use it. This is the biggest issue in this code.

Explain what the union-find data structure does, generally, in plain english. If you want, explain it in terms of contacts, that's fine. Update the docstrings of WUF, addRoot, getRoot, getRootGroups, and union to reflect this in plain english. (for example, ContactTracer, add_contact, get_network, get_all_networks, merge_networks with associated docstrings)

For example:

ContactTracer uses contacts between pairs of people to track networks of people who might have infected one another. If two people are in contact, ContactTracer always puts them in the same network. Further, if Alice is in contact with Bob who is in contact with Carol, all three are in one network (and so on). If there is no chain of contacts between two people, they are in different networks. ContactTracer can dynamically update with new people and contacts in real time. It can never remove people or contacts.

Note that this description suggests a useful new function -- check_same_network(a: person_id, b: person_id): bool

It's too long

This is a really long file to do not that much. For comparison, here is a related question asked by another python user.

The docstring format might be a little long combined with the types, honestly (due to the parameter documentation). It's okay. Too much documentation is better than too little.

However, remove:

  • __eq__ (won't be used and doesn't do what you want)
  • __iter__ (won't be used)
  • __len__ (won't be used)
  • __getitem__ (won't be used)


Path compression is the main reason people talk about union-find structures. It really improves performance, and it's about three lines of code. Add it. If you add path compression, I would argue you no longer really need the weighting half, but that's up to you.

if key not in rootGroups[root]: is slow (a linear check through every item in the list). Omit the check entirely. It's not needed.

Follow standard python style

Read PEP8. Prefer snake_case to camelCase in python.


Your security question is nonsense. Security doesn't usually make sense without talking about a specific external interface, like a web interface. User-supplied data of some kind. Also, expectations as to what security this class should offer.


You've mentioned that you're using this for contact tracing, I assume for COVID-19 infection modelling. This is not a very good contract tracing method. You may be getting false positives by not taking time into account.

For example, suppose Alice talks to Bob, and then Bob talks to Carol, and then Carol contracts a disease from Dave. Then Alice and Bob are disease-free unless they talk to Carol again. But your algorithm lumps Alice, Bob, Carol, and Dave all into one union, and will probably report that all four could be infected.

  • \$\begingroup\$ Thank you, I very much appreciate this response. A lot of this is really helpful and stuff I'll look to implement. As for the last bit about correctness, this was for an introductory programming class final (it was a prereq) so I wasn't super concerned with accuracy to that degree. \$\endgroup\$
    – Gabe Ron
    Commented Jun 24, 2022 at 8:18

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.