[Codeforces] 1519C. Berland Regional 풀이

Problem

Key Idea

• 각 university 별로 따로 계산이 가능하다.
• 한 university에 학생 수가 \(sz\) 라면, \(1 \le k \le sz\) 인 k값에만 영향이 있다.

첫 번째 예시.

7
1 2 1 2 1 2 1
6 8 3 1 5 1 5

각 university 별로 list에 담고, sorting하면

univ[1]: 3 5 5 6
univ[2]: 1 1 8

\(univ[1]\) 계산: \(k = 3\) 일때 6이 제외되면서, \(ans[3]\)에만 6이 빠진다.

ans: 0 19 19 13 19 0 0 0

즉, \(sz \mod k\) 개 만큼 앞에서부터(작은것부터) 빠진다.

1. university별로 sort.
2. \(univ[i] (1 \le i \le n)\) 배열을 누적합 배열로 변환
3. \(univ\) 배열을 순회하면서, 각 \(k\)값 \((1 \le k \le sz)\)에 대해 \((총합 - univ[i][sz\%k-1])\) 을 \(ans[k]\)에 더해준다.

• Time: \(O(n{\log}n)\)

Implementation

/**
 * author: jooncco
 * written: 2021. 5. 8. Sat. 00:22:56 [UTC+9]
 **/

#include <bits/stdc++.h>
using namespace std;

#define FAST_IO ios_base::sync_with_stdio(0),cin.tie(0)
typedef long long ll;
typedef vector<ll> vl;

int t,n;
ll u[200010], s;

int main() {

    FAST_IO;
    cin >> t;
    while (t--) {
        cin >> n;
        for (int i=0; i < n; ++i) cin >> u[i];
        vector<vl> univ(n+1,vl());
        ll ans[200010]= {0};
        for (int i=0; i < n; ++i) {
            cin >> s;
            univ[u[i]].push_back(s);
        }
        for (int i=1; i <= n; ++i) {
            int sz= univ[i].size();
            if (sz) {
                sort(univ[i].begin(),univ[i].end());
                for (int j=1; j < sz; ++j) univ[i][j] += univ[i][j-1];
                for (int k=1; k <= sz; ++k) {
                    ans[k] += univ[i][sz-1];
                    if (sz%k > 0) ans[k] -= univ[i][sz%k-1];
                }
            }
        }
        for (int i=1; i <= n; ++i) {
            if (i > 1) cout << " ";
            cout << ans[i];
        }
        cout << "\n";
    }
}