[LeetCode] 2615 - Sum of Distances

문제

• 링크: https://leetcode.com/problems/sum-of-distances
• 난이도: Medium
• 태그: 배열, 해시 테이블, 누적합
• 결과: Time: 53 ms (88.00%), Space: 109.52 MB (87.75%)

풀이

class Solution {
public:
    vector<long long> distance(vector<int>& nums) {
        int n = nums.size();
        unordered_map<int, vector<int>> indexMap;
        for (int i = 0; i < n; ++i)
        {
            indexMap[nums[i]].push_back(i);
        }

        vector<long long> v(n);

        for (const auto& p : indexMap)
        {
            const auto& group = p.second;
            long long sum = accumulate(group.begin(), group.end(), 0LL);
            long long prefixSum = 0;
            for (int i = 0; i < group.size(); ++i)
            {
                v[group[i]] = sum - prefixSum * 2 + group[i] * (2 * i - group.size());
                prefixSum += group[i];
            }
        }

        return v;
    }
};

오랜만에 푸는 Daily Question이다. 일단 해시 맵으로 인덱스를 그룹짓는 것 까지는 생각했는데, 누적합과 결합할 생각은 못해서 Editorial을 보고 풀었다.

막혔던 지점

결국 $|x-i|+|x-j|+\dots+|x-m|$을 구해야하는데, 절댓값이 들어가다보니 상수시간에 구하는 법을 생각해내지 못했다.

해결 방법

누적합 아이디어를 가져오고, 왼쪽에 있는 요소와 오른쪽에 있는 요소를 나누어서 거리를 구하면 된다.

가장 먼저, 같은 수 끼리는 인덱스를 기록해둔다. 이 때, 차례대로 훑으며 기록하기 때문에 이 ‘그룹’의 인덱스는 오름차순으로 정렬되어 있다.

구하고자 하는 $\text{res}[a_i]=\sum^{m-1}_{j=0}{

a_i-a_j

}$에서 그룹별 인덱스가 정렬되어있음을 이용하여 왼쪽과 오른쪽으로 나누고 합칠 수 있다. 그룹 내에 있는 모든 인덱스의 합을 $S$, 누적합을 $P_i=\sum^{i-1}_{j=0}a_j$라고 하자.

왼쪽은 아래와 같이 구할 수 있다.

\[\sum^{i-1}_{j=0}{(a_i-a_j)}=\sum^{i-1}_{j=0}{a_j}=i\times a-P_i\]

오른쪽은 아래와 같이 구할 수 있다.

\[\sum^{m-1}_{j=i+1}{(a_j-a_i)} =\sum^{m-1}_{j=i+1}{a_j}-\sum^{m-1}_{j=i+1}{a_i} =(S-P_i-a_i)-a_i\times(m-i-1)\]

$\sum^{m-1}_{j=i+1}{a_j}$은 모든 인덱스의 합 $S$에서 $i$까지의 누적합을 뺀 것과 같으므로 $(S-P_i-a_i)$가 된다.

이제 왼쪽과 오른쪽을 합치면 최종적으로 아래와 같다.

\(\text{res}[a_i]=(i\times a-P_i)+((S-P_i-a_i)-a_i\times(m-i-1))=S-2P_i+a_i\times(2i-m)\) 이제 위 수식을 그대로 코드로 옮기면 된다.

배운 점과 후기

std::accmumulate를 처음 알았다. 그리고 누적합 아이디어를 다른 알고리즘에 응용하는 풀이를 익혔다. 이러한 응용에 더 익숙해져야할 것 같다.