[LeetCode 4] Median of Two Sorted Arrays 해설

leetcode.com/problems/median-of-two-sorted-arrays/

Median of Two Sorted Arrays - LeetCode

 

1. 분류

Queue

 

2. 풀이

두 int 배열의 평균이 아니라 중앙값을 구하는 문제이다.

 

두 배열을 합친 크기가 최대 2,000 밖에 안 되기 때문에 그냥 한 배열에 넣고 정렬한 뒤 중앙값 구해도 풀 수 있다. (...)

 

나는 두 우선순위 큐를 이용한 방법으로 해결했다.

대부분의 언어에 구현되어 있는 우선순위 큐는 힙 기반이며 최소힙 / 최대힙을 사용할 수 있다.

 

먼저, 최소힙과 최대힙 1개씩을 만들어둔다.

그리고 아래 조건에 맞춰서 힙에 숫자를 넣는다.

1. 큰 수는 최소 힙에, 작은 수는 최대 힙에 넣을 것이다.
2. 최소 힙의 top은 최대 힙의 top보다 무조건 크다.
   • 최소 힙의 모든 원소는 최대 힙의 모든 원소보다 큰 값이다.
3. 최소 힙의 크기는 최대 힙의 크기보다 크거나 같다.

예시를 보면 이해가 빠르다.

 

1. [1, 2, 3, 4, 5]
   • 최소 힙: [3, 4, 5]   (최대 힙의 최고 큰 숫자인 2보다 모두 크다)
     • top: 3
   • 최대 힙: [1, 2]
     • top: 2
2. [1, 2, 3, 4, 5, 6]
   • 최소 힙: [4, 5, 6]   (최대 힙의 최고 큰 숫자인 3보다 모두 크다)
     • top: 4
   • 최대 힙: [1, 2, 3]
     • top: 3

 

위에서 정의한 힙 제약조건이 있기 때문에 아래 가정은 무조건 참이 된다.

• 최소 힙 크기 = 최대 힙 크기인 경우 각 힙의 top을 더해 반으로 나눈 값 = 중앙값
• 최소 힙 크기 > 최대 힙 크기인 경우 최소 힙의 top = 중앙값

 

3. 소스코드

3-1) 두 배열 합치고 정렬하는 단순한 방법

import java.util.Arrays;
import java.util.stream.IntStream;

class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {

        int[] concat = IntStream.concat(Arrays.stream(nums1), Arrays.stream(nums2))
                                .sorted()
                                .toArray();

        if (concat.length == 1) {
            return concat[0];
        }

        if (concat.length % 2 == 0) {
            return (double)(concat[concat.length / 2 - 1] + concat[concat.length / 2]) / 2;
        }

        return concat[concat.length / 2];
    }
}

 

3-2) 힙을 이용하는 방법

import java.util.Comparator;
import java.util.PriorityQueue;

class Solution {

    PriorityQueue<Integer> min = new PriorityQueue<>();
    PriorityQueue<Integer> max = new PriorityQueue<>(Comparator.reverseOrder());

    public double findMedianSortedArrays(int[] nums1, int[] nums2) {

        enqueue(nums1);
        enqueue(nums2);

        if (min.size() > max.size()) {
            return (double) min.peek();
        }

        return (((double) min.peek()) + ((double) max.peek())) / 2f;
    }

    private void enqueue(int[] nums) {
        for (int num : nums) {
            if (min.size() <= max.size()) {
                min.add(num);
            } else {
                max.add(num);
            }

            if (min.size() == 0 || max.size() == 0) {
                continue;
            }

            if (min.peek() <= max.peek()) {
                int a = min.poll();
                int b = max.poll();

                max.add(a);
                min.add(b);
            }
        }
    }
}

 

- 테스트 코드 -

import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.MatcherAssert.assertThat;

import java.util.stream.Stream;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;

class SolutionTest {

    private final Solution code = new Solution();

    @MethodSource("testcase")
    @ParameterizedTest
    void test(int[] nums1, int[] nums2, double output) {

        var result = code.findMedianSortedArrays(nums1, nums2);
        System.out.println("result: " + result);
        System.out.println("output: " + output);

        assertThat(result, is(output));
    }

    private static Stream<Arguments> testcase() {
        return Stream.of(
            Arguments.of(new int[] {1,3}, new int[] {2}, 2.00000),
            Arguments.of(new int[] {0,0}, new int[] {0,0}, 0.00000),
            Arguments.of(new int[] {}, new int[] {1}, 1.00000),
            Arguments.of(new int[] {2}, new int[] {}, 2.00000)
        );
    }
}