Redis Disjoint Set

redis_disjoint_set.py

import random

class RedisDisjointSet(object):
    """
    https://en.wikipedia.org/wiki/Disjoint-set_data_structure
    
    UTF-8 based persistent disjoint set structure that just assumes there's a local redis
    """
    
    def __init__(self, seed=None):
        if seed is None:
            self.seed = ''.join(random.choices(string.ascii_uppercase + string.digits, k=N))
        else:
            self.seed = seed

        self.r = redis.Redis(host='localhost', port=6379, db=0)
        
        self.redis_key = f'disjoint_set:{self.seed}'
        self.l = f'{self.redis_key}:leader'
        self.g = f'{self.redis_key}:group'

    def leader(self, k=None, v=None):
        if k is not None and v is not None:
            self.r.hset(self.l, k, v)
        elif k is not None:
            
            v = self.r.hget(self.l, k)
            if v: return v.decode('utf-8')
        else:
            return {
                k.decode('utf-8'): v.decode('utf-8')
                for k,v in self.r.hgetall(self.l).items()
            }
        
    def group(self):
        leaders = self.leader().values()
        return {
            l: [v.decode('utf-8') for v in self.r.smembers(f"{self.g}:{l}")]
            for l in leaders
        }
        
    def add(self, a: str, b: str):
        leadera = self.leader(a)
        leaderb = self.leader(b)
        groupa = f'{self.g}:{leadera}'
        groupb = f'{self.g}:{leaderb}'
        if leadera is not None:
            if leaderb is not None:
                if leadera == leaderb: return # nothing to do
                card_groupa = self.r.scard(groupa)
                card_groupb = self.r.scard(groupb)
                if card_groupa < card_groupb:
                    a, leadera, groupa, b, leaderb, groupb = b, leaderb, groupb, a, leadera, groupa
                    
                self.r.sunionstore(groupa, groupa, groupb)

                groupb_items = self.r.smembers(groupb)
                self.r.delete(groupb)
                
                for k in groupb_items:
                    self.leader(k, leadera)
            else:
                self.r.sadd(groupa, b)
                self.leader(b, leadera)
        else:
            if leaderb is not None:
                self.r.sadd(groupb, a)
                self.leader(a, leaderb)
            else:
                self.leader(b, a)
                self.leader(a, a)
                self.r.sadd(f'{self.g}:{a}', a)
                self.r.sadd(f'{self.g}:{a}', b)                
                
    def clear(self):
        for k in self.r.keys(f'{self.redis_key}*'):
            self.r.delete(k)
            
if __name__ == "__main__":
    ds = RedisDisjointSet('test')

    ds.add('1', '3')
    ds.add('1', '4')
    ds.add('3', '5')
    ds.add('4', '3')
    ds.add('5', '1')
    ds.add('6', '7')
    ds.add('8', '4')

    print(ds.leader())
    print(ds.group())

    ds2 = RedisDisjointSet('test')
    print(ds2.leader())

    ds.clear()
    print(ds.leader())