AtCoder ABC 107 D번 Median of Medians

작성일 | In

문제 링크

  • https://atcoder.jp/contests/abc107/tasks/arc101_b

문제 출처

  • AtCoder ABC 107 D번
  • AtCoder ARC 191 B번

사용 알고리즘

  • 펜윅 트리
  • 파라메트릭 서치

시간복잡도

  • $O(N log N log 10^9)$

풀이

모든 $(l, r)$쌍 $(1 ≤ l ≤ r ≤ N)$에 대해, $m_{l, r}$의 값들을 모아 새로 만든 수열의 중앙값을 찾는 문제입니다.

길이가 $N$인 수열 $A$의 중앙값이 $m$이라면 수열 $A$에는 $x$보다 크거나 같은 수가 $\lceil {\frac{N}{2}} \rceil$개 이상 있어야 하고, $x$는 이 조건을 만족하는 가장 큰 수여야 합니다.

수열 $A$에서 어떤 수 $x$보다 크거나 같은 수를 1로 바꾸고, $x$보다 작은 수를 0으로 바꿔준 새로운 수열 $B$를 생각해봅시다.
만약 수열 $B$의 정답이 1이라면 수열 $A$에서의 정답은 $x$ 이상이고, 수열 $B$의 정답이 0이라면 수열 $A$에서의 정답은 $x$ 미만이 됩니다.
이런 느낌으로, 어떤 수 $x$를 고정시켜서 $m_{l, r} ≥ x$를 만족하는 $(l, r)$쌍의 개수를 세는 문제로 바꿔서 parametric search를 시도할 수 있습니다.

위에서 언급했듯이, $m_{l, r} ≥ x$가 성립하기 위해서는 $x$이상인 수의 개수가 $x$미만인 수의 개수보다 많거나 같아야합니다. $x$이상인 수를 1, $x$미만인 수를 -1로 바꿔주면, $A_l + A_{l+1} + A_{l+2} + … + A_r ≥ 0$으로 표현할 수 있습니다.

조건을 만족하는 $(l, r)$쌍의 개수를 구하기 위해서, prefix sum을 활용해봅시다. $A_l + A_{l+1} + A_{l+2} + … + A_r ≥ 0$은 $S_{l-1} ≤ S_r$과 같습니다. $S_{l-1} ≤ S_r$을 만족하는 $(l, r)$쌍의 개수를 세는 문제는 inversion counting이라는 이름으로 널리 알려져 있습니다.

prefix sum을 fenwick tree로 구현해주면서 parametric search를 해주면 $O(N log N log 10^9)$에 문제를 해결할 수 있습니다.

전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

ll n;
ll arr[101010];
ll tree[202020];

void update(int x){
	while(x < 202020){
		tree[x]++;
		x += x & -x;
	}
}

int query(int x){
	int ret = 0;
	while(x){
		ret += tree[x];
		x -= x & -x;
	}
	return ret;
}

bool chk(int m){
	memset(tree, 0, sizeof tree);
	ll res = 0, t = n + 1;
	update(t);
	for(int i=1; i<=n; i++){
		t += arr[i] >= m ? 1 : -1;
		res += query(t);
		update(t);
	}
	return res >= ((n*n+n)/2+1)/2;
}

int main(){
    ios_base::sync_with_stdio(0); cin.tie(0);
    cin >> n;
	for(int i=1; i<=n; i++) cin >> arr[i];

    int l = 0, r = 1e9;
	while(l < r){
		int m = l + r + 1 >> 1;
		if(chk(m)) l = m;
		else r = m - 1;
	}
	cout << l;
}