UnionFind in python (enhanced version: strings and tuples are allowed for elements)

Introduction

Although it is the famous Union Find, there are few elements of the set that can handle strings and tuples, so I will write an article. I haven't encountered many useful situations, so I'm not sure if it will be useful, but if you are interested, please take a look.

Reference article

https://note.nkmk.me/python-union-find/ https://qiita.com/white1107/items/52fd4149bb1846862e38

Code introduction

I'll leave it to other articles to explain what UnionFind looks like, and I'll show you the code right away.

UnionFind.py


import collections
class UnionFind():
    def __init__(self):
        '''
with unionfind path compression,Elements can be tuples or strings,No number of elements specified at the beginning
        '''
        self.parents = dict()                                   #{element:number,}
        self.membersf = collections.defaultdict(lambda : set()) #Tuples and strings are allowed in set
        self.rootsf = set()                                     #Tuples and strings are acceptable
        self.d = dict()
        self.d_inv = dict()
        self.cnt = 0
    
    def dictf(self,x):
        if x in self.d:
            return self.d[x]
        else:
            self.cnt += 1
            self.d[x] = self.cnt
            self.parents[x] = self.cnt
            self.d_inv[self.cnt] = x
            self.membersf[x].add(x)
            self.rootsf.add(x)
            return self.d[x]

    def find(self, x):
        self.dictf(x)
        if self.parents[x] == self.dictf(x):
            return x
        else:
            self.parents[x] = self.d[self.find(self.d_inv[self.parents[x]])]
            return self.d_inv[self.parents[x]]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if self.parents[x] > self.parents[y]:
            x, y = y, x
        if x == y:
            return
        for i in list(self.membersf[y]):
            self.membersf[x].add(i)
        self.membersf[y] = set()
        self.rootsf.remove(y)
        self.parents[y] = self.d[x]

    def size(self, x):#Number of elements in the set containing x
        return len(self.membersf[self.find(x)])

    def same(self, x, y):#Judgment whether they belong to the same set
        return self.find(x) == self.find(y)

    def members(self, x):#Elements of the set containing x
        return self.membersf[self.find(x)]

    def roots(self):#Root element
        return self.rootsf

    def group_count(self):#Number of roots
        return len(self.rootsf)

    def all_group_members(self):#Roots and their elements
        return {r: self.membersf[r] for r in list(self.rootsf)}

if __name__=="__main__": #For debugging
    uf_s = UnionFind()
    uf_s.union('A', 'D')
    uf_s.union('D', 'C')
    uf_s.union('E', 'B')
    print(uf_s.members('D'))
    print(uf_s.members('E'))
    uf_s.union('A', 'B')
    print(uf_s.members('E'))
    print(uf_s.same('A','F'))
    print(uf_s.members('F'))
    print(uf_s.members('G'))

    print()
    uf_t = UnionFind()
    uf_t.union((2,3),(3,4))
    uf_t.union((4,3),(5,4))
    uf_t.union((1,3),(3,4))
    print(uf_t.members((2,3)))
    print(uf_t.roots())
    print(uf_t.group_count())
    print(uf_t.all_group_members())

    print()
    uf_ts = UnionFind()
    uf_ts.union((2,3),1)
    uf_ts.union((4,3),'A')
    uf_ts.union((1,3),(3,4))
    uf_ts.union((1,3),'A')
    print(uf_ts.members((2,3)))
    print(uf_ts.roots())
    print(uf_ts.all_group_members())
    print(uf_ts.size('A'))

I've attached a nice debug code below.

Explanation

How to use

Basically, you can use Union Find, which you often see, in the same way. find (x): Find the root of x union (x, y): union sets size (x): Outputs the size of the set containing x elements same (x, y): Determine if they belong to the same set members (x): Output a set containing x roots (): get all roots group_count (): Output the number of sets all_group_members (): Get all roots and their elements

Ingenuity and advantages

Anyway, by making full use of the dictionary, we manage the set by compressing the route while taking the correspondence between the element and the element number. Therefore, any element that can be used as a key in a dictionary can be treated in the same way even if it has a different type (string, tuple, integer, etc.). In addition, many union finds seen elsewhere specify the number of elements first, but in this code, every time an element arrives, it is checked whether it is the first time, and if it is the first time, a number that has not been used until now is assigned, so the element You do not need to specify the number. (Addition: It can also be used when coordinate compression is required or when you want to do 2D union find.)

Weaknesses

Compared to ordinary union find, it consumes a lot of memory because there are many parts managed by dictionaries. In addition, there is a lot of data exchange inside unionfind, and it takes a lot of calculation time. Example: https://atcoder.jp/contests/practice2/tasks/practice2_a I didn't do TLE, but it seems that it may not pass if it gets heavier than 1500ms.

in conclusion

Thank you for reading. I myself am a beginner in competitive programming, so please point out any mistakes. (In addition, if you have any problems with this code, please comment on the link etc.)

Recommended Posts

UnionFind in python (enhanced version: strings and tuples are allowed for elements)
Search for strings in Python
Tips for those who are wondering how to use is and == in Python
What are python tuples and * args after all?
Rock-paper-scissors poi in Python for beginners (answers and explanations)
I tried pipenv and asdf for Python version control
Preferences for playing Wave in Python PyAudio and PortAudio
Get the size (number of elements) of UnionFind in Python
Recursively search for files and directories in Python and output
Compare strings in Python
Reverse strings in Python
Seeking a unified way to wait and get for state changes in Selenium for Python elements
Make sure all the elements in the list are the same in Python
For those who are having trouble drawing graphs in python
List method argument information for classes and modules in Python
How to swap elements in an array in Python, and how to reverse an array.
Sort and output the elements in the list as elements and multiples in Python.
Installation procedure for Python and Ansible with a specific version
Tips for coding short and easy to read in Python
Useful tricks related to list and for statements in Python
Introduction to Effectiveness Verification Chapters 4 and 5 are written in Python
Problems and solutions when asked for MySQL db in Python 3
3-3, Python strings and character codes
Search for strings in files
Techniques for sorting in Python
Stack and Queue in Python
Python list and tuples and commas
Unittest and CI in Python
Getting list elements in Python
About "for _ in range ():" in python
tesseract-OCR for Python [Japanese version]
Python> list> append () and extend ()> append: list is added | extend: list elements are added | + = to add list
Regular expressions that are easy and solid to learn in Python
Search for character strings in files [Comparison between Bash and PowerShell]
Create a CGH for branching a laser in Python (laser and SLM required)