넘치게 채우기
[LeetCode] 3321. Find X-Sum of All K-Long Subarrays II 본문
Leetcode - Find X-Sum of All K-Long Subarrays II
문제 유형: 슬라이딩 윈도우, 집합, 해시맵
문제 난이도: Hard
문제
You are given an array nums of n integers and two integers k and x.
The x-sum of an array is calculated by the following procedure:
- Count the occurrences of all elements in the array.
- Keep only the occurrences of the top x most frequent elements. If two elements have the same number of occurrences, the element with the bigger value is considered more frequent.
- Calculate the sum of the resulting array.
Note that if an array has less than x distinct elements, its x-sum is the sum of the array.
Return an integer array answer of length n - k + 1 where answer[i] is the x-sum of the subarray nums[i..i + k - 1].
n개 길이의 정수배열 nums를 받고, 정수 k와 x가 주어진다.
배열의 x-sum은 다음과 같은 방식으로 구해진다:
- 배열의 요소별 모든 개수를 구한다.
- 가장 많은 빈도 x개를 골라 (빈도 * 수)를 누적한 값을 x-sum이라 한다. 같은 빈도라면 더 큰수를 고른다.
- 빈도가 x개 이하이면, 배열의 수를 모두 더한다.
answer[i] = nums[i..i+k-1]의 x-sum인 n-k+1짜리 길이의 배열 answer를 반환하시오.
풀이
전용 자료구조를 만드는 것이 편하다.
Helper라고 하자. 아래와 같은 필드들을 가진다고 가정하자:
- 빈도계산을 위한 occ 해시맵
- x
- <빈도, 숫자>를 담은 레드블랙트리 기반의 set인 large, small. large에는 top x개의 요소를 유지하고, small에는 나머지 후보들을 저장한다.
- xsum의 결과인 result.
그 뒤, insert(), remove(), get()을 구현한다.
insert()에는, 이전에 빈도가 있으면, 내부적으로 그 값을 지우게 하고, 빈도를 1 증가시키고, 내부적으로 새 레코드<새 빈도, 숫자>를 추가시키도록 한다.
remove()는 내부적으로 이전 레코드를 지우고, 빈도를 1 내리고, 여전히 빈도가 1 이상이라면 새 레코드<기존빈도-1, 숫자>를 추가시키도록 한다.
get()은 result를 반환하도록 한다.
이제, 내부적인 레코드 관리 로직인 internalInsert(), internalRemove()를 구현한다.
internalInsert()에서는 large 셋의 크기가 x 미만이거나, 이번에 삽입하는 레코드 p가 가장 작은 요소보다 크다면, large셋에 추가하며,
result에 새로 값을 추가한다. 그게 아니라면, small셋에 추가한다.
interanlRemove()에서는 p가 large의 가장 작은 레코드 이상이라면, large레코드에서 찾아서 없애고, result의 값에 pop된 값을 지우며 새로 업데이트한다. 그러면서 small에 있는 가장 큰 값을 large로 끌어올리며 다시 result에 더하며 업데이트한다.
p가 small에 있다면, small로부터만 지워주면 된다.
메인함수에서는, 매번 insert해주면 되고, i가 k이상이면 remove(), k-1이상이면 get()을 호출하여 그 값을 각각 저장하면된다.
코드
C++
class Helper {
public:
Helper(int x) {
this->x=x;
this->result=0;
}
void insert(int num) {
if(occ[num]) {
internalRemove({occ[num], num});
}
++occ[num];
internalInsert({occ[num], num});
}
void remove(int num) {
internalRemove({occ[num], num});
--occ[num];
if(occ[num]) {
internalInsert({occ[num], num});
}
}
long long get() { return result; }
private:
void internalInsert(const pair<int, int>& p) {
if(large.size() < x || p > *large.begin()) {
result += 1LL * p.first * p.second;
large.insert(p);
if(large.size() > x) {
result -= 1LL * (large.begin()->first) * (large.begin()->second);
auto transfer = *large.begin();
large.erase(transfer);
small.insert(transfer);
}
} else {
small.insert(p);
}
}
void internalRemove(const pair<int, int> p) {
if(p >= *large.begin()) {
result -= 1LL * p.first * p.second;
large.erase(p);
if(!small.empty()) {
result += 1LL * (small.rbegin()->first) * (small.rbegin()->second);
auto transfer = *small.rbegin();
small.erase(transfer);
large.insert(transfer);
}
} else {
small.erase(p);
}
}
int x;
long long result;
set<pair<int, int>> large, small;
unordered_map<int, int> occ;
};
class Solution {
public:
vector<long long> findXSum(vector<int>& nums, int k, int x) {
Helper helper(x);
vector<long long> ans;
for(int i = 0; i < nums.size(); ++i) {
helper.insert(nums[i]);
if(i >= k) {
helper.remove(nums[i-k]);
}
if(i >= k-1) {
ans.emplace_back(helper.get());
}
}
return ans;
}
};