[LeetCode # 315] Count of Smaller Numbers After Self

정수형 배열 nums 가 주어졌을 때, 각 인덱스 별로 자기보다 오른쪽에 있는 수 중 자기보다 작은 수들의 개수를 구하는 문제

 

예를 들어, 배열 nums 가 [5, 2, 6, 1] 과 같이 주어졌을 때 정답은 [2, 1, 1, 0] 이다. 가장 오른쪽에 있는 '1' 은 자기보다 오른쪽에 있는 값이 없기 때문에 무조건 0이 되고, '6' 과 '2' 는 오른쪽에 '1' 이, '5'는 오른쪽에 '2', '1' 이 있어 각각 1 과 2 를 정답으로 갖는다.

 

이 문제를 풀기 위한 방법으로 두 가지 방법을 사용했다. 첫 번째 방법은 Merge Sort 를 활용한 방법이었고 두 번째는 Binary Index Tree 의 일종인 Fenwick Tree 를 활용한 방법이다.

 

우선, Merge Sort 를 활용한 방법은 아이디어가 전혀 떠오르지 않아 아래 Reference 를 참고했다. Merge Sort 는 기본적으로 Divide & Conquer 방식으로 정렬을 수행하는데, 이 때 매 단계 별로 시작 위치 ~ 끝 위치 중 가운데 위치를 선정하고, 시작 위치 ~ 가운데 위치 / 가운데 위치 + 1 ~ 끝 위치까지 정렬을 먼저 한 후 시작 위치 ~ 끝 위치의 정렬을 수행한다. 정렬을 할 때, 시작 위치 ~ 가운데 위치 / 가운데 위치 + 1 ~ 끝 위치는 각각 정렬이 된 상태이므로 (아래 단계에서 이미 정렬을 하고 난 상태) 두 배열을 (위치 별로 나눈 것을 두 배열이라 표현) 합치는 과정에서 두 배열 중 작은 값을 먼저 배치하는 방식으로 정렬을 한다. 이를 간단히 코드로 표현하면 다음과 같다.

 

void merge (int[] arr, int[] tmp, int st, int ed) {
	for (int i = st; i <= ed; i++)
    	tmp[i] = arr[i];
	int mid = (st + ed) / 2;
    int i = st, j = mid + 1, k = st;
    while (i <= mid && j <= ed) {
    	if (tmp[i] <= tmp[j])
        	arr[k++] = tmp[i++];
        else
        	arr[k++] = tmp[j++];
    }
    while (i <= mid)
    	arr[k++] = tmp[i++];
    while (j <= ed)
    	arr[k++] = tmp[j++];
}

 

이 문제를 해결하기 위해, Merge Sort 에 추가적으로 도입한 아이디어는 두 배열을 하나로 합치는 과정에서 뒤 배열의 값들이 앞 배열보다 먼저 배치될 경우 앞 배열의 값들은 원래 위치에서 자기보다 오른쪽에 먼저 배치된 값들이 존재한다는 것이다. 즉, 위의 merge 소스 코드에서 앞 배열과 뒤 배열의 값들을 서로 비교해 합치는 과정에서 뒤 배열의 값이 더 크다면 그러한 개수만큼 저장을 하고 앞 배열의 값을 배치할 때 해당 값이 원래 위치했던 인덱스의 결과를 증가시키면 되는 것이다.

 

이러한 해결 방법으로 작성한 코드는 아래와 같다.

 

struct node {
    int idx, val;
};

class Solution {
public:
    void merge(vector<int>& ans, vector<node>& numVector, vector<node>& copyVector, int st, int mid, int ed) {
        for (int i = st; i <= ed; i++) {
            copyVector[i] = numVector[i];
        }
        int i = st, j = mid + 1, k = st, cnt = 0;

        while (i <= mid && j <= ed) {
            if (copyVector[i].val <= copyVector[j].val) {
                ans[copyVector[i].idx] += cnt;
                numVector[k++] = copyVector[i++];
            }
            else {
                cnt++;
                numVector[k++] = copyVector[j++];
            }
        }

        while (i <= mid) {
            ans[copyVector[i].idx] += cnt;
            numVector[k++] = copyVector[i++];
        }
        while (j <= ed) {
            numVector[k++] = copyVector[j++];
        }
    }

    void mergeSort(vector<int>& ans, vector<node>& numVector, vector<node>& copyVector, int st, int ed) {
        if (st >= ed)
            return;
        int mid = (st + ed) / 2;

        mergeSort(ans, numVector, copyVector, st, mid);
        mergeSort(ans, numVector, copyVector,mid + 1, ed);

        merge(ans, numVector, copyVector, st, mid, ed);
    }

    vector<int> countSmaller(vector<int>& nums) {
        int n = nums.size();
        vector<int> ans(n, 0);
        vector<node> numVector(n);
        vector<node> copyVector(n);
        for (int i = 0; i < n; i++) {
            numVector[i].idx = i;
            numVector[i].val = nums[i];
        }
        mergeSort(ans, numVector, copyVector, 0, n - 1);
        return ans;
    }
};

 

위 방식으로 문제를 해결할 수는 있었지만, 시간복잡도는 O(NlogN) 으로 (N 은 주어진 nums 배열의 크기, 최대 10만) 수행시간이 생각보다 길어 다른 방식을 고민하게 되었다. 이 문제에서 nums 배열의 크기는 최대 10만이고, 배열에 저장되는 값들은 -1만 ~ 1만 의 값을 가진다. 즉 O(NlogN) 보다는, O(NlogM) 이 된다면 (이 때 M 은 배열에 저장되는 값의 범위, 이 문제에서는 최대 2만) 수행시간을 조금 더 단축할 수 있을 것으로 생각했다.

 

이 때, logN 을 logM 으로 바꿀 수 있는 방법으로 부분합을 생각했고, 부분합을 업데이트하는 횟수가 많아 세그먼트 트리보다는 펜윅 트리를 생각했다. 펜윅 트리에 대한 자세한 설명은 아래 Reference 를 참고하면 되고, 간단하게 얘기하면 값들의 업데이트가 많을 때 부분합을 구하고 업데이트하는 것의 시간복잡도가 모두 O(logN) 인 자료구조이다.

 

펜윅 트리는 기본적으로 인덱스를 사용해 인덱스 ? ~ ? 까지의 합을 저장하는 방식으로 사용하는데, 이 문제에서는 인덱스 대신 nums 배열의 값들을 사용했고 값이 -1만 ~ 1만의 범위를 가지기 때문에, Offset 으로 10001 을 더해서 사용했다. 트리에 저장되는 값들은 해당 값이 등장한 횟수이다.

 

펜윅 트리를 사용해 작성한 코드는 아래와 같다. 수행시간도 생각한 대로 Merge Sort 를 사용한 방식보다 빨랐고, 펜윅 트리의 Update 나 Sum 을 구하는 코드가 모두 간단하기 때문에 코드도 간단하게 작성할 수 있었다.

 

class Solution {
public:
    const int OFFSET = 10001;
    const int SIZE = 20010;
    vector<int> countSmaller(vector<int>& nums) {
        int n = nums.size();
        vector<int> tree(SIZE, 0);
        vector<int> ans(n, 0);

        for (int i = n - 1; i >= 0; i--) {
            int cnt = 0;
            for (int j = nums[i] + OFFSET - 1; j > 0; j -= (j & -j))
                ans[i] += tree[j];
            for (int j = nums[i] + OFFSET; j < SIZE; j += (j & -j))
                tree[j]++;
        }

        return ans;
    }
};

 

위 코드를 로직에만 집중하도록 간단하게 모듈화를 조금 더 하자면 다음과 같이 작성할 수 있을 것 같다.

 

class Solution {
public:
    const int OFFSET = 10001;
    const int SIZE = 20010;
    void update(vector<int>& tree, int idx) {
        while (idx < SIZE) {
            tree[idx]++;
            idx += (idx & -idx);
        }
    }
    int sum(vector<int>& tree, int idx) {
        int sum = 0;
        while (idx > 0) {
            sum += tree[idx];
            idx -= (idx & -idx);
        }
        return sum;
    }
    vector<int> countSmaller(vector<int>& nums) {
        int n = nums.size();
        vector<int> tree(SIZE, 0);
        vector<int> ans(n, 0);

        for (int i = n - 1; i >= 0; i--) {
            int cnt = 0;
            ans[i] = sum(tree, nums[i] + OFFSET - 1);
            update(tree, nums[i] + OFFSET);
        }

        return ans;
    }
};

 

- Reference

[GeeksforGeeks - using MergeSort] https://www.geeksforgeeks.org/count-of-smaller-elements-on-right-side-of-each-element-in-an-array-using-merge-sort/

Count of smaller elements on right side of each element in an Array using Merge sort - GeeksforGeeks

[Baekjoon - Fenwick Tree] https://www.acmicpc.net/blog/view/21

펜윅 트리 (바이너리 인덱스 트리)