염소줄서기 문제 풀이
2060 염소줄서기 풀이 및 코드
개요
오랜만에 다이아 문제 풀이를 써보려고 한다. 이 문제는 내가 옛날에 북마크 해뒀던 문젠데 북마크에서 거의 1년간 썩어가고 있어서 속상해서 선택했다. 문제 이해가 어렵지 않아서 선택한 것도 있다. 다이아 문제를 하루에 하나씩 풀면 엄청난 도움이 되겠지만 블로그 풀이를 안보면 하루에 하나는 무슨 일주일에 하나도 힘들것이다. 그래도 1일 1플레는 도전을 해봐야겠다.
문제 이해
염소들 번호가 있는데 밥먹는 순서는 다음과 같다. 염소번호를 이진수로 나타냈을 때
- 1의 개수가 적을 수록 먼저먹는다.
- 1의 개수가 같다면 그냥 이진수 크기가 더 작은 염소가 먼저 먹는다.
이진수로 A부터 B까지의 염소가 있을 때, k번째 밥을 먹는 염소는 몇 번 친구인가?
Naive한 풀이
당연히 정렬을 생각할 수 있다. 하지만 1의 개수가 최대 31개, 따라서 integer형 범위를 모두 포함한다면 약 20억이다. $A = 0$, $B = 2^{32} - 1$ 이라면 절대 순수하게 정렬할 수는 없을 것이다.
손으로 하는 풀이 내가 내 눈과 손으로 답을 구할 때는 먼저 1의 개수가 작은 것부터 보면서 counting을 했다. 1의 개수가 1개인것이 $a_1$개 2개 인것이 $a_2$ 개 .. 이런식으로 생각하면서 k를 줄여나갔다.
풀이
세 단계로 풀이를 세웠다.
- 1의 개수가 $i$개인 염소의 수를 dp table에 전처리 해둔다.($0 \leq i \leq 31$)
- dp table의 누적합을 이용해 k번째 염소가 몇개의 1이 있는지 확인한다.
- c개의 1이 있는 염소들 중 k - 어쩌구 번째 친구를 구한다.
1단계
rnk(i, X) : X보다 작은 1의 개수가 i개인 수의 개수
A와 B가 고정이기 때문에 rnk(i, B) - rnk(i, A-1)을 하면 i별로 1의 개수가 i개인 염소의 수를 구할 수 있다.
rnk(i, X)
예를 들어 1001011보다 작은 1의 개수가 5개인 수들을 구해보자! 아래 예시에서 첫번째 요소를 볼때 0이 오는지 1이 오는지에 따라 경우의 수를 세줄 수 있다.
즉, 1을 만났을 때 $\binom{이후 자리수}{남은 1의 수}$ 를 더해주면 된다.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
ll rnk(int cnt, string s){
ll ret = 0, one = 0;
for(int i=0;i<s.length();i++){
if (s[i] == '1') one++;
}
for(int i=0;i<s.length();i++){
if ((s.length() - i == one) || one == 0 || cnt == 0){
ret += binom(one, cnt);
break;
}
if(s[i] == '1'){
ret += binom(s.length()-i-1, cnt);
one -= 1;
cnt -= 1;
}
}
return ret;
}
이렇게 rnk함수를 구할 수 있었다. binom함수는 2차원 dp로 구현했다. 처음에 one == 0, cnt == 0이 되는 경우를 처리하지 않아서 계속 틀렸다.
2단계
누적합을 이용하면 $O(31)$ 이내에 간단하게 몇개의 비트를 쓰는지 확인할 수 있다.
1
2
3
4
5
6
7
8
9
10
11
for(int i=0;i<32;i++){
dp[i] = ((i == 0) ? 0 : dp[i-1]) + dp[i];
}
int c = 0;
ll kk = 0;
for(int i=0;i<32;i++){
if(dp[i] >= k) break;
c = i;
kk = dp[i];
}
k -= kk;
k번째 수가 c개의 비트를 가지고 있다는 뜻이고 kk는 지금까지 건너뛴 염소들의 수다. 이제 A와 B사이에 있는 c개의 비트를 가진 수들 중에서 k(-= kk)번째 수를 구하면 된다.
3단계
사실 이 과정에서 막혔다… c개의 비트를 가진 수들 중에서 k번째 수를 어떻게 구할까..? 근데 잘 생각해보니 그냥 아까의 rnk함수를 이용해서 이분탐색을 하면 됐다. 왜냐면 x보다 작은 1의 개수가 c인 수는 정렬되어있기 때문이다. 그리고 그것이 변하는 순간은 무조건 c개의 비트를 가지고 있는것이 명확하기 때문에 제일 먼저 k가 나오는 순간을 포착하면 답이 된다.
전체 코드
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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
#include <bits/stdc++.h>
#define fast_io cin.tie(NULL); ios_base::sync_with_stdio(false);
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef tuple<int, int, int> tiii;
typedef tuple<ll, ll, ll> tlll;
#define xx first
#define yy second
ll comb[32][32];
ll binom(int n, int k){
if(n < k) return 0;
if(k == 0 || n == k) return 1;
ll &ret = comb[n][k];
if(ret != -1) return ret;
ret = binom(n-1, k) + binom(n-1, k-1);
return ret;
}
ll rnk(int cnt, string s){
ll ret = 0, one = 0;
for(int i=0;i<s.length();i++){
if (s[i] == '1') one++;
}
for(int i=0;i<s.length();i++){
if ((s.length() - i == one) || one == 0 || cnt == 0){
ret += binom(one, cnt);
break;
}
if(s[i] == '1'){
ret += binom(s.length()-i-1, cnt);
one -= 1;
cnt -= 1;
}
}
return ret;
}
string to_binary(ll x){
if(!x) return "0";
string ret = "";
while(x){
ret += (x % 2) + '0';
x /= 2;
}
reverse(ret.begin(), ret.end());
return ret;
}
ll to_ll(string s){
if(s == "-1") return -1;
ll ret = 0;
ll j = 1;
for(int i=s.length()-1;i>=0;i--){
ret += j * (s[i] - '0');
j *= 2;
}
return ret;
}
int main(){
fast_io
string A, B;
cin >> A >> B;
ll k; cin >> k;
assert (to_ll(B) - to_ll(A) + 1 >= k);
memset(comb, -1, sizeof(comb));
if(A != "0"){
for (int i = A.length() - 1; i >= 0; i--){
if (A[i] == '0') A[i] = '1';
else{
A[i] = '0';
if(i == 0){
if(A.length() != 1) A.erase(A.begin());
}
break;
}
}
}else{
A = "-1";
}
vector<ll> dp(32, 0);
for(int i=0;i<32;i++){
dp[i] = rnk(i, B) - (A == "-1" ? 0 : rnk(i, A));
}
for(int i=0;i<32;i++){
dp[i] = ((i == 0) ? 0 : dp[i-1]) + dp[i];
}
int c = 0;
ll kk = 0;
for(int i=0;i<32;i++){
c = i;
if(dp[i] >= k) break;
kk = dp[i];
}
k -= kk;
ll lo = to_ll(A), hi = to_ll(B);
while(lo + 1 < hi){
ll mid = (lo + hi) >> 1;
if(rnk(c, to_binary(mid)) - (A == "-1" ? 0 : rnk(c, A))< k) lo = mid;
else hi = mid;
}
cout << to_binary(hi);
}
마무리
솔직히 엄청 어려운 알고리즘은 하나도 안들어간다. 진법 변환, 자릿수 dp, 이항계수, 누적합 등 매우 쉬운 알고리즘 들이다. 항상 보면 이런 자잘한 것들이 모여서 어려운 문제를 만들어낸다..ㅠㅠ 이 문제의 포인트는 누적합 논리인것 같다. 그리고 누적합에서는 이분탐색이 가능하다는 것도 잘 이용한 좋은 문제라고 생각한다.