[leetcode] 49. Group Anagrams 파이썬 코드

 

https://leetcode.com/problems/group-anagrams/description/

class Solution(object):
    def groupAnagrams(self, strs):
        """
        :type strs: List[str]
        :rtype: List[List[str]]
        """
        result = [[] for i in range(len(strs))]
        group_dict_pos = []

        for i in range(len(strs)):
            vocab = strs[i]
            vocab_dict = dict()
            for alphabet in vocab:
                if alphabet in vocab_dict:
                    vocab_dict[alphabet] += 1
                else:
                    vocab_dict[alphabet] = 1

            if vocab_dict in group_dict_pos:
                idx = group_dict_pos.index(vocab_dict)
                result[idx].append(vocab)
            else:
                group_dict_pos.append(vocab_dict)
                idx = group_dict_pos.index(vocab_dict)
                result[idx].append(vocab)

        if [] in result:
            final_result = []
            for i in range(len(result)):
                if not result[i]:
                    pass
                else:
                    final_result.append(result[i])
            
            return final_result
        else:
            return result

 

나의 코드는 효율적이지 않았다.

 

 

그럼 어떻게 효율적으로 풀 수 있을까?

 

일단, 두 단어가 서로 anagram임을 알 수 있는 방법은 뭘까?

(1) 단어를 swap 해서 같은 단어인지 확인한다

(2) 알파벳 숫자를 카운트한다

(3) 정렬을 해본다 

 

정렬을 하면, time complexity는 O(n logn)일 것이고 n은 input string의 평균적인 길이다.

주어지는 단어의 전체 개수가 m개라고 하면, 전체적인 시간 복잡도는 O(m * n logn)이다.

 

이것보다 더 좋게 만들 수 있을까?

 

딕셔너리의 value 값에 리스트를 집어 넣으려면 defaultdict를 사용하면 된다는 걸 알게 되었다. 

 

from collections import defaultdict

strs = ["eat","tea","tan","ate","nat","bat"]

res = defaultdict(list)

for s in strs:
    count = [0] * 26

    for c in s:
        count[ord(c) - ord('a')] += 1

    res[tuple(count)].append(s)
    # res[count].append(s) # 딕셔너리의 key 자리에 리스트는 올 수 없다


print(res)
print(res.values())

 

 

참고

https://www.youtube.com/watch?v=vzdNOK2oB2E&t=9s