BOJ 33802 Difference Maximization

문제 링크입니다.

문제 요약

$1$ 이상 $M$ 이하의 정수 $N$개 중 일부의 값을 알고 있을 때, 나머지를 $1$ 이상 $M$ 이하로 적절히 배정해 임의의 두 원소의 차들의 합이 최대가 되도록 하세요.

풀이

실력을 알고 있는 친구들 사이에서 얻을 수 있는 친밀도를 먼저 계산합시다. 친밀도의 합이라는 건 결국 각 $1 \le i < M$마다 (실력이 $i$ 이하인 친구 수)×(실력이 $i+1$ 이상인 친구 수)의 합을 구하는 것과 같습니다. 각 실력마다 친구 수를 세서 누적 합을 이용해 이를 ${\cal O}(N+M)$에 계산할 수 있습니다.

실력을 알 수 없는 친구들에게 각각 얼마의 실력을 줄지 생각해 봅시다. 뭔가 어중간한 실력을 주는 것보단 극단적인 값을 주는 게 좀더 이득일 것 같아요. 간격을 최대한 벌리는 게 목적이니까…

어떤 친구가 다른 친구의 실력에 따라 얻는 친밀도는 절댓값 함수의 형태로 나타납니다. 한 친구를 어떤 실력에 추가했을 때 얻게 되는 총 친밀도는 이 절댓값 함수의 합으로 나타낼 수 있습니다. 그러면 그래프의 형태가 쭉 내려가다가 어느 순간부터 다시 올라가는, 볼록함수 형태임을 알 수 있습니다.

우리는 이 함수의 최댓값이 필요합니다. 이는 항상 학생의 실력이 $1$이나 $M$ 중 하나일 때 나타납니다. 어느 쪽에서 나타날지는 모르니 친구를 추가할 때마다 $1$과 $M$에 친구를 추가했을 때 얻게 될 친밀도를 업데이트하면서 그때그때 더 좋아 보이는 쪽을 그리디하게 골라주면 됩니다. 그리디가 작동함을 믿을 수 없다면 $1$과 $M$에 추가할 친구 수를 달리해가면서 ${\cal O}(N)$에 브루트포싱해도 될 겁니다.

위의 그리디 전략이 성립하는 이유는 다음과 같습니다. 일반성을 잃지 않고 $1$과 $M$ 중 $1$을 선택하는 쪽이 현재 더 좋다고 합시다. 그런데도 $M$을 추가한 경우에, 두 가지 케이스가 있을 수 있습니다.

1. 이후로 모두 $M$을 선택한 경우 : $M$을 선택할 때마다 $1$을 선택했을 때 얻을 수 있는 친밀도는 늘어납니다. 반면 $M$을 선택했을 때 얻는 친밀도는 그대로이므로, 마지막에 $1$을 선택하는 쪽이 반드시 최종 친밀도가 더 커집니다. 따라서 이런 경우가 최선의 해일 수 없습니다.
2. 이후로 $1$을 선택한 적이 있는 경우 : 고르는 순서에 따라 답이 변하는 게 아니기 때문에 지금 $1$을 골라도 같은 최선의 해를 얻을 수 있습니다. 위의 1번 케이스에서 현재 이득인 쪽을 적어도 한 번은 골라야 최적해를 얻을 수 있음을 확인했으므로, 두 경우를 조합하면 매 순간마다 이득인 쪽을 고르는 최적해가 존재함을 알 수 있습니다.

따라서 이상의 내용을 구현하면 총 시간복잡도 ${\cal O}(N+M)$에 문제를 해결할 수 있습니다.

코드

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
using pii = pair<int, int>;
const ll mod = 1e9+7;

ll A[1010101], chmin, chmax;

int main(){
	cin.tie(0);cout.tie(0);ios::sync_with_stdio(false);
	int n, m; cin >> n >> m;
	int cnt = 0; ll ans = 0;
	for(int i = 1; i <= n; i++){
		int a; cin >> a;
		if(!a){
			cnt++; continue;
		}
		chmin += a-1;
		chmax += m-a;
		A[a]++;
	}
	ll sum = 0;
	for(int i = 1; i < m; i++){
		sum += A[i];
		ans += sum*(n-cnt-sum);
	}
	while(cnt--){
		if(chmin < chmax){
			chmin += m-1;
			ans += chmax;
		}
		else{
			chmax += m-1;
			ans += chmin;
		}
	}
	cout << ans;
}