Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

Python으로 각 쿼리를 포함하는 가장 짧은 구간의 크기 찾기

문제 이해하기

구간(interval) 목록이 주어진다고 가정해 보겠습니다. 여기서 intervals[i]는 쌍 (left_i, right_i)으로 표현되며, i번째 구간이 left_i에서 시작하여 right_i에서 끝난다는 의미입니다(양쪽 경계 모두 포함). 또한 queries라는 배열도 함께 주어집니다.

j번째 쿼리에 대한 답은 left_i <= queries[j] <= right_i 조건을 만족하는, 즉 해당 쿼리 값을 포함하는 구간 중에서 크기가 가장 작은 구간의 길이입니다. 만약 조건을 만족하는 구간이 하나도 없다면 -1을 반환합니다. 최종적으로는 모든 쿼리에 대한 답을 담고 있는 배열을 반환해야 합니다.

예시 살펴보기

입력이 다음과 같다고 해보겠습니다.

  • intervals = [[2,5],[3,5],[4,7],[5,5]]
  • queries = [3,4,5,6]

이 경우 출력은 [3, 3, 1, 4]가 됩니다. 각 쿼리는 다음과 같이 처리됩니다.

  • query = 3: 3을 포함하는 가장 작은 구간은 [3,5]이므로, 5 − 3 + 1 = 3
  • query = 4: 4를 포함하는 가장 작은 구간은 [3,5]이므로, 5 − 3 + 1 = 3
  • query = 5: 5를 포함하는 가장 작은 구간은 [5,5]이므로, 5 − 5 + 1 = 1
  • query = 6: 6을 포함하는 가장 작은 구간은 [4,7]이므로, 7 − 4 + 1 = 4

알고리즘 접근 방법

이 문제를 효율적으로 해결하려면 정렬최소 힙(min-heap)을 활용할 수 있습니다. 단계별 과정은 다음과 같습니다.

  1. intervals 목록을 내림차순으로 정렬합니다.
  2. 힙으로 사용할 빈 리스트 h와 결과를 저장할 맵 res를 준비합니다.
  3. 쿼리를 오름차순으로 정렬한 뒤, 각 쿼리 q에 대해 다음을 수행합니다.
    • intervals가 비어 있지 않고 마지막 구간의 시작 시점이 q 이하일 동안, 구간을 하나씩 꺼냅니다. 이때 구간의 끝 시점(j)이 q 이상이라면, 즉 현재 쿼리를 포함할 가능성이 있다면 (길이 j−i+1, 끝 시점 j) 쌍을 힙 h에 삽입합니다.
    • 힙 h가 비어 있지 않고 힙 최상단 구간의 끝 시점이 q보다 작으면, 더 이상 유효하지 않으므로 제거합니다(pop).
    • 힙 h가 비어 있지 않다면 res[q]에 최상단 요소의 길이를 저장하고, 비어 있다면 -1을 저장합니다.
  4. 마지막으로 원래 쿼리 순서대로 res[q] 값들을 리스트로 만들어 반환합니다.

핵심 아이디어는 쿼리를 정렬된 순서로 처리하면서, 시작 지점이 현재 쿼리 이전에 있는 구간들을 힙에 미리 넣어 두고, 범위를 벗어난 구간은 힙에서 제거하는 방식입니다. 이렇게 하면 매 쿼리마다 유효한 구간 중 가장 짧은 것을 O(log n) 시간에 얻을 수 있습니다.

Python 구현 예제

아래 코드를 통해 실제 구현을 확인해 보겠습니다.

import heapq

def solve(intervals, queries):
    intervals = sorted(intervals)[::-1]
    h = []
    res = {}
    for q in sorted(queries):
        while intervals and intervals[-1][0] <= q:
            i, j = intervals.pop()
            if j >= q:
                heapq.heappush(h, [j - i + 1, j])
        while h and h[0][1] < q:
            heapq.heappop(h)
        res[q] = h[0][0] if h else -1
    return [res[q] for q in queries]

intervals = [[2,5],[3,5],[4,7],[5,5]]
queries = [3,4,5,6]
print(solve(intervals, queries))

입력

[[2,5],[3,5],[4,7],[5,5]], [3,4,5,6]

출력

[3, 3, 1, 4]