프로그래머스 가사 검색 코드 및 해설 (파이썬)

https://programmers.co.kr/learn/courses/30/lessons/60060

코딩테스트 연습 - 가사 검색

 

먼저 단어의 길이가 key, 해당 길이의 단어들의 리스트를 value로 갖는 word_dict를 생성합니다. 같은 방식으로 단어의 앞뒤를 뒤집은 reversed_word_dict도 생성합니다.
matching_dict라는 딕셔너리를 만들어 key는 하나의 쿼리, value는 해당 쿼리에 매칭되는 단어의 개수를 저장합니다.
이후 queries를 하나씩 꺼내면서 현재의 쿼리가 ?로만 이루어져 있으면 해당 쿼리 길이의 단어들의 개수를 matching_dict에 저장합니다.
쿼리가 ?를 접미사로 가지면, 해당 쿼리 길이의 단어들을 array, target은 ?를 제외한 알파벳으로 두고 이진 탐색을 실행합니다. 이진 탐색 함수는 정렬된 배열에서 특정 수의 개수 구하기의 함수를 약간만 수정했습니다. 'ab' < 'ac' 등 string의 알파벳 정렬 순서로 <, >, == 등의 연산이 가능함을 이용했습니다.
쿼리가 ?를 접두사로 가지면, 쿼리의 앞뒤를 뒤집고, 단어들도 뒤집은 형태를 활용하여 ?를 접미사로 가질 때와 동일하게 이진탐색을 실행했습니다.
모든 쿼리에 대해 매칭 가능한 단어의 개수를 구한 뒤, 각 쿼리당 단어 개수를 리스트로 만들어 최종 정답으로 반환했습니다.

 

https://codlingual.tistory.com/189

이진 탐색

 

from collections import defaultdict

# 이진 탐색 
def binary_search(array, suffix_idx, target, start, end, result):
    if start > end:
        return len(result) 

    mid = (start + end) // 2 

    # 중간값이 타겟인 경우 
    if array[mid][:suffix_idx] == target:
        # 현재 인덱스를 result 집합에 추가 
        result.add(mid)
        # 현재 인덱스보다 작은 쪽, 큰 쪽 모두 가보기 
        binary_search(array, suffix_idx, target, start, mid - 1, result)
        binary_search(array, suffix_idx, target, mid + 1, end, result)

    # 중간값이 타겟보다 큰 경우 
    elif target < array[mid][:suffix_idx]:
        binary_search(array, suffix_idx, target, start, mid - 1, result)

    # 중간값이 타겟보다 작은 경우 
    else:
        binary_search(array, suffix_idx, target, mid + 1, end, result)

    return len(result) 

def solution(words, queries):
    # key: 단어 길이, value: 해당 길이의 단어들 
    word_dict = defaultdict(list)
    # key: 단어 길이, value: 해당 길이의 단어들의 앞뒤를 뒤집은 형태 
    reversed_word_dict = defaultdict(list)
    for word in words:
        word_dict[len(word)].append(word)
        reversed_word_dict[len(word)].append(word[::-1])

    # value를 알파벳 순으로 정렬 
    word_dict = {key: sorted(words) for key, words in word_dict.items()}
    reversed_word_dict = {key: sorted(words) for key, words in reversed_word_dict.items()}
    
    matching_dict = defaultdict(int)
    
    for q in set(queries):
        # query가 ?로만 이루어진 경우 
        if set(q) == set('?'):
            if len(q) not in word_dict:
                matching_dict[q] = 0
                continue        
            matching_dict[q] = len(word_dict[len(q)])
            
        # query의 접미사가 ?인 경우 
        elif not q.startswith('?'):
            if len(q) not in word_dict:
                matching_dict[q] = 0
                continue 
            array = word_dict[len(q)]
            suffix_idx = q.index('?')
            target = q[:suffix_idx]
            matching_dict[q] = binary_search(array, suffix_idx, target, 0, len(array)-1, result=set())
            
        # query의 접두사가 ?인 경우 
        # query의 앞뒤를 뒤집어 접미사인 경우와 동일하게 탐색 
        else:
            if len(q) not in reversed_word_dict:
                matching_dict[q] = 0
                continue 
            array = reversed_word_dict[len(q)]
            prefix_idx = q[::-1].index('?')
            target = q[::-1][:prefix_idx]
            matching_dict[q] = binary_search(array, prefix_idx, target, 0, len(array)-1, result=set())
    
    # 중복된 query가 있을 수 있으므로 
    answer = []
    for query in queries:
        answer.append(matching_dict[query])
    return answer