1 분 소요

문제

내 풀이

from collections import defaultdict
from typing import List

def shortestDistanceColor(colors: List[int], queries: List[List[int]]) -> List[int]:
    answer = []
    n = len(colors)
    # dist[i][1] : i번째 인덱스에서 1까지 가장 가까운 거리
    dist = [[0]*4 for _ in range(n)]   # nx4 행렬
    
    # dist 채우기
    for i in range(n):
        for k in range(1,4):
            # dist[i][k] 구하기
            flag = False
            for d in range(n):  # i로부터 거리
                if i-d >= 0 and colors[i-d] == k:
                    dist[i][k] = d
                    flag = True
                    break

                if i+d < n and colors[i+d] == k:
                    dist[i][k] = d
                    flag = True
                    break

            if not flag:
                dist[i][k] = -1

    # queries 답 찾기 answer[i] = dist[queries[i][0]][queries[i][1]]
    for i in range(len(queries)):
        answer.append(dist[queries[i][0]][queries[i][1]])

    return answer

colors = [1,2]
queries = [[0,3]]
answer = shortestDistanceColor(colors, queries)
print(answer)
  • 걍 bruteforce로 풀었음
  • 당연하게도 시간 초과남
  • 시간 초과 이유 : c에 해당하지 않는 값도 going through 하느라 시간 낭비
  1. initializing 3 lists : one for each color
  2. colors 배열 순회 → putting each index into its respective color list (최적화)
  3. queries 배열 순회 → c 컬러에 해당하는 리스트만 순회하면 됨.
  4. 최적화) 3에서 c컬러에 해당하는 리스트 전체를 조회할 필요 x ⇒ binary search 적용
def shortestDistanceColor(colors: List[int], queries: List[List[int]]) -> List[int]:
    hashmap = defaultdict(list)
    for i,c in enumerate(colors):
        hashmap[c].append(i)

    query_results = []
    for i, (target, color) in enumerate(queries):
        if color not in hashmap:
            query_results.append(-1)
            continue

        index_list = hashmap[color]
        # 왜..? target에서 가장 가까운 color의 위치 찾는거 아닌가 -> list에서 target과 가장 가까운 값 찾으려고
        insert = bisect.bisect_left(index_list, target)

        # compare the index on the left and right of insert
        # make sure it will not fall out of the index_list
        left_nearest = abs(index_list[max(insert - 1 , 0)] - target)
        right_nearest = abs(index_list[min(insert, len(index_list) - 1)] - target)
        query_results.append(min(left_nearest, right_nearest))

    return query_results

참고