1202. Smallest String With Swaps

[zeikar]

Problem link

https://leetcode.com/problems/smallest-string-with-swaps/

Problem Summary

문자열과 스왑이 가능한 인덱스가 주어질 때 스왑을 해서 문자열을 사전 순으로 최소가 되게 만드는 문제.

Solution

조금만 생각해보면 Disjoint-set 으로 풀 수 있다.
스왑을 하는 인덱스가 연결이 되어 있다면 얼마든지 스왑을 통해 연결된 인덱스의 최소로 바꿀 수 있다는 점을 이용한 것이다.
예를 들어 [0, 1], [1, 2]가 스왑이 가능하면 여러 번 스왑으로 2번 인덱스가 0번 인덱스로 넘어갈 수 있다.

연결이 되었는지 여부는 Disjoint-set으로 간단하게 구현할 수 있다.
단, find, union 연산은 대충 구현하면 O(N)이라 최적화를 좀 해야 시간 안에 들어온다. 아래 링크의 path compression 과 union by rank를 둘 다 구현.

연결이 되어 있는 인덱스에 대해 heap (pq)로 작은 문자부터 뽑아다가 결과 문자열을 만들면 된다. 시간 복잡도는 O(N logN)

참고: https://leetcode.com/explore/learn/card/graph/618/disjoint-set/3844/

Source Code

import heapq
from collections import defaultdict
from typing import List


class Solution:
    def smallestStringWithSwaps(self, s: str, pairs: List[List[int]]) -> str:
        def find(x):
            if x != parent[x]:
                parent[x] = find(parent[x])
            return parent[x]

        def union(x, y):
            px, py = find(x), find(y)
            if px != py:
                if rank[px] > rank[py]:
                    parent[py] = px
                else:
                    parent[px] = py
                    if rank[px] == rank[py]:
                        rank[py] += 1

        parent = [i for i in range(len(s))]
        rank = [1] * len(s)
        for x, y in pairs:
            union(x, y)

        m = defaultdict(list)
        for i in range(len(s)):
            m[find(i)].append(s[i])

        for i in m:
            heapq.heapify(m[i])

        res = ""
        for i in range(len(s)):
            res += heapq.heappop(m[find(i)])
        return res