leetcode medium : Exam Room

https://leetcode.com/problems/exam-room/description/

풀기 전

빈공간이 제일 큰 부분을 빠르게 찾아서 삽입해야됨

1 <= n <= 10^9 이고 At most 104 calls will be made 인걸 보니 O(N)으로는 좀 힘들어보이고 O(log N)알고리즘이 필요해보임.

‘제일 큰걸 빠르게 찾는다’ → Max heap 연상

뺄 때도 O(log N)이 가능한지 고민해보았으나 average speed는 좀 줄일 수 있어도 worst case에서의 속도는 줄이기 힘들다고 판단. 빼는건 naive하게 O(N)으로 구현함.

극단적인 케이스

1. seat만 m번 ⇒ O(m log n)
2. seat m/2번하고 leave m/2번 ⇒ O(mn)

코드 설명

class ExamRoom:
    def push(self, begin, end):
        if end < begin or end > self.N - 1 or begin < 0:
            return

        dist = (end - begin) // 2
        
        if begin == 0:
            dist = end
        if end == self.N - 1:
            dist = self.N - 1 - begin
        
        heapq.heappush(self.heap, [-dist, begin, end])

    def __init__(self, n: int):
        self.heap = []
        self.N = n
        self.push(0, n - 1)

    def seat(self) -> int:
        dist, begin, end = heapq.heappop(self.heap)
        dist *= -1

        if begin == 0:
            new_seat_pos = begin
        elif end == self.N - 1:
            new_seat_pos = end
        else:
            new_seat_pos = begin + dist
        
        self.push(begin, new_seat_pos-1)
        self.push(new_seat_pos+1, end)

        return new_seat_pos

    def leave(self, p: int) -> None:
        before_p = None
        after_p = None
        
        for i in range(len(self.heap)):
            dist, begin, end = self.heap[i]
            dist *= -1

            if end == p - 1:
                before_p = self.heap[i]
            if begin == p + 1:
                after_p = self.heap[i]

        begin = before_p[1] if before_p else p
        end = after_p[2] if after_p else p

        if before_p:
            self.heap.remove(before_p)
        if after_p:
            self.heap.remove(after_p)

        self.push(begin, end)
        heapq.heapify(self.heap)

초기화

def __init__(self, n: int):
	  self.heap = []
	  self.N = n
	  self.push(0, n - 1)

힙을 초기화하고 인덱스 0부터 n-1까지 빈공간이 있음을 표시함

빈공간 표시 함수 (heap에 push)

def push(self, begin, end):
    if end < begin or end > self.N - 1 or begin < 0:
        return

    dist = (end - begin) // 2
    
    if begin == 0:
        dist = end
    if end == self.N - 1:
        dist = self.N - 1 - begin
    
    heapq.heappush(self.heap, [-dist, begin, end])

1. 파라미터 체크
2. 이 빈공간에 앉게 되면 가장 가까운 사람과의 거리가 얼마나 되는지 dist에 계산
   1. 엣지 케이스 : 끝쪽 빈공간의 경우(begin == 0 or end == N - 1) 맨 끝에 앉아버리면 dist를 최대화할수 있으므로 빈공간//2로 계산하면 안된다.
3. 빈공간을 힙에 푸시
   1. -dist 를 넣는 이유는 파이썬 힙이 기본적으로 min heap이기 때문ㅇ
   2. begin을 두 번째 인자로 넣는 이유는 dist가 같은 빈공간이 여러개있을때 앞쪽 빈공간에 seat하기 위해

def seat(self) -> int:
    dist, begin, end = heapq.heappop(self.heap)
    dist *= -1

    if begin == 0:
        new_seat_pos = begin
    elif end == self.N - 1:
        new_seat_pos = end
    else:
        new_seat_pos = begin + dist
    
    self.push(begin, new_seat_pos-1)
    self.push(new_seat_pos+1, end)

    return new_seat_pos

1. 힙에서 distance to closest person이 maximize될 수 있는 빈공간 꺼내기
2. 어느 인덱스에 앉아야되는지 계산 (엣지케이스 고려해서 if문 사용)
3. 앉은 인덱스 기준으로 양쪽 빈공간을 힙에 푸시

def leave(self, p: int) -> None:
    before_p = None
    after_p = None
    
    for i in range(len(self.heap)):
        dist, begin, end = self.heap[i]
        dist *= -1

        if end == p - 1:
            before_p = self.heap[i]
        if begin == p + 1:
            after_p = self.heap[i]

    begin = before_p[1] if before_p else p
    end = after_p[2] if after_p else p

    if before_p:
        self.heap.remove(before_p)
    if after_p:
        self.heap.remove(after_p)

    self.push(begin, end)
    heapq.heapify(self.heap)

before_p : p 왼쪽의 빈공간

after_p : p 오른쪽의 빈공간

힙을 linear scan해서 p에 붙어있는 빈공간이 있으면 찾아서 힙에서 제거

p를 뺐을 때 새로 생기는 빈공간을 heap에 푸시

분석

시간복잡도는 seat가 O(log N), leave가 O(N)

heapify가 O(N)인 이유

공간복잡도는 약 n/2명이 seat하게 될경우 모든 사람이 평균적으로 사이에 1개의 빈공간을 끼고 있을 수 있어서 빈공간도 약 n/2개가 될수 있음 → O(N)

솔루션 학습

leave를 O(log N)으로 처리할수있을지 다른 사람들의 솔루션을 살펴봤다.

솔루션 1

class ExamRoom(object):

    def __init__(self, N):
        """
        :type N: int
        """
        self.N = N
        self.heap = []
        self.avail_first = {}
        self.avail_last = {}
        self.put_segment(0, self.N - 1)

    def put_segment(self, first, last):

        if first == 0 or last == self.N - 1:
            priority = last - first
        else:
            priority = (last - first) // 2

        segment = [-priority, first, last, True]

        self.avail_first[first] = segment
        self.avail_last[last] = segment

        heappush(self.heap, segment)

    def seat(self):
        """
        :rtype: int
        """
        while True:
            _, first, last, is_valid = heappop(self.heap)

            if is_valid:
                del self.avail_first[first]
                del self.avail_last[last]
                break

        if first == 0:
            ret = 0
            if first != last:
                self.put_segment(first + 1, last)

        elif last == self.N - 1:
            ret = last
            if first != last:
                self.put_segment(first, last - 1)

        else:
            ret = first + (last - first) // 2

            if ret > first:
                self.put_segment(first, ret - 1)

            if ret < last:
                self.put_segment(ret + 1, last)

        return ret

    def leave(self, p):
        """
        :type p: int
        :rtype: void
        """
        first = p
        last = p

        left = p - 1
        right = p + 1

        if left >= 0 and left in self.avail_last:
            segment_left = self.avail_last.pop(left)
            segment_left[3] = False
            first = segment_left[1]

        if right < self.N and right in self.avail_first:
            segment_right = self.avail_first.pop(right)
            segment_right[3] = False
            last = segment_right[2]

        self.put_segment(first, last)

leave는 O(log N)이 맞는데 seat부분에서

while True:
    _, first, last, is_valid = heappop(self.heap)

    if is_valid:
        del self.avail_first[first]
        del self.avail_last[last]
        break

invalid한 빈공간은 스킵하면서 pop하는 부분이있다.

생각해보자 n/2명이 seat하고 n/2-1명이 leave했다가 다시 1명이 seat하면?

마지막에 seat할때 n에 비례하는 invalid한 빈공간들이 pop되기 때문에 O(n log n)의 시간복잡도가 된다.

물론 연산의 실행횟수는 10000개 이하로 제한되어있으므로 연산의 실행횟수가 m이라고 할 때 O(m log m)이 소요될거고 10000이면 그리 크지 않아서 무리가 되진 않을것이다.

사실 생각해보면 m이 작아서 O(N)알고리즘이어도 그냥 통과되지 않을까 싶긴하다.

솔루션 2

struct IntV
{
    static int N;
    int l, r;
    IntV(int l, int r) : l(l), r(r) {}
    int getD() const
    {
        if(l == 0) return r;
        if(r == N - 1) return N - 1 - l;
        if(r < l) return -1;
        else return (r - l) / 2;
    }
    int getP() const
    {
        if(l == 0) return 0;
        if(r == N - 1) return N - 1;
        else return l + (r - l) / 2;
    }
    bool operator < (const IntV& a) const
    {
        int d1 = getD(), d2 = a.getD();
        if(d1 != d2) return d1 > d2;
        return l < a.l;
    }
};

int IntV :: N = 0;

class ExamRoom {
public:
    set<IntV> tab;
    map<int, int> l2r, r2l;
    
    ExamRoom(int N)
    {
        IntV :: N = N;
        tab.clear(); l2r.clear(); r2l.clear();
        tab.insert(IntV(0, N - 1));
        l2r[0] = N - 1;
        r2l[N - 1] = 0;
    }
    
    int seat()
    {
        IntV cur = *(tab.begin());
        tab.erase(tab.begin());
        
        int pos = cur.getP();
        tab.insert(IntV(cur.l, pos - 1));
        tab.insert(IntV(pos + 1, cur.r));
        l2r[cur.l] = pos - 1; r2l[pos - 1] = cur.l;
        l2r[pos + 1] = cur.r; r2l[cur.r] = pos + 1;
        return pos;
    }
    
    void leave(int p)
    {
        int l = r2l[p - 1], r = l2r[p + 1];
        tab.erase(IntV(l, p - 1));
        tab.erase(IntV(p + 1, r));
        tab.insert(IntV(l, r));
        l2r[l] = r; r2l[r] = l;
        r2l.erase(p - 1); l2r.erase(p + 1);
    }
};

내 솔루션에서 leave를 O(log N)으로 만들기 위해서는 다음 조건을 만족해야 했다.

• p를 뺌으로써 양쪽에 invalid한 공간 두 개가 남는데 이를 O(log N) 타임에 뺄 수 있어야 한다. → heap이라 중간요소 제거 불가능

def leave(self, p: int) -> None:
    before_p = None
    after_p = None
    
    **for i in range(len(self.heap)):
        dist, begin, end = self.heap[i]
        dist *= -1

        if end == p - 1:
            before_p = self.heap[i]
        if begin == p + 1:
            after_p = self.heap[i]**

    begin = before_p[1] if before_p else p
    end = after_p[2] if after_p else p

    if before_p:
        self.heap.remove(before_p)
    if after_p:
        self.heap.remove(after_p)

    self.push(begin, end)
    **heapq.heapify(self.heap)**

굵은 글씨 친 부분이 O(N)을 만드는 부분이다.

이 문제를 해결하려면 힙이 아닌 다른 자료구조를 써야했다.

힙 = 삽입과 삭제가 O(log N), 대신 삭제는 가장 크거나 작은 요소 한정

우리에게 필요한거 = 임의의 요소를 삽입하고 삭제하는대 O(log N)이 걸리면서 가장 크거나 작은 요소를 O(1)만에 찾을 수 있는 자료구조 = AVL tree

AVL tree는 기본적으로 BST기 때문에 정렬되어있어 가장 크거나 작은 요소를 빠르게 찾을 수 있고 삽입/삭제 역시 logTime에 가능하다. 이 솔루션의 작성자는 C++ stl의 set가 red black tree라는 것에 착안하여 다음과 같이 가장 빈공간이 큰 요소를 빠르게 빼내었다.

IntV cur = *(tab.begin());

그리고 l2r, r21이라는 2개의 map을 만들어 특정인덱스로부터 시작하는 빈공간이 어디서 끝나는지 O(1)타임에 알수있도록 하였다.

매우 smart한 풀이인듯