[BOJ] 9527 - 1의 개수 세기

https://www.acmicpc.net/problem/9527

BOJ - 1의 개수 세기

문제 유형: 누적합, 수학, 비트마스킹

문제 난이도: Gold II

시간 제한: 1초

메모리 제한: 128MB

 

문제

두 자연수 A, B가 주어졌을 때, A ≤ x ≤ B를 만족하는 모든 x에 대해 x를 이진수로 표현했을 때 1의 개수의 합을 구하는 프로그램을 작성하시오.

즉, f(x) = x를 이진수로 표현 했을 때 1의 개수라고 정의하고, 아래 식의 결과를 구하자.

 ∑x=ABf(x)\[\sum_{x=A}^{B}{f(x)}\] 

 

입력

첫 줄에 두 자연수 A, B가 주어진다. (1 ≤ A ≤ B ≤ 10^16)

 

출력

1의 개수를 세어 출력한다.

 

풀이

어떤 수 2^i-1에 대한 점화식은 다음과 같다:

 

pows[i]는 (1 << i)-1까지의 숫자에서 1의 개수를 계산해놓은 누적합이다.

 

i비트의 경우마다 다음과 같다:

0: 0개

1: 1개

 

1까지 1개

 

00: 0개

01: 1개

10: 1개

11: 2개

 

3까지 4개

 

000: 0개

001: 1개

010: 1개

011: 2개

100: 1개

101: 2개

110: 2개

111: 3개

 

7까지 12개.

 

pows[i]를 0부터 2^i-1까지의 누적합이라고 하자.

pows[i] = pows[i-1] * 2 + 2^(i+1)이라는,

즉 이전 비트범위까지의 값이 두 번 있고, 개수만큼 더해주면 된다.

 

이제 각 MSB부터 읽으면서, 만약 1이라면 

부분합 pows[i-1]을 누적한다.

그러고 1xxxxx의 시작, 즉 10000...000부터 현재 수까지의 차 + 1, 즉 닫힌 구간의 개수를 누적해준다.

이는 위의 예시에서처럼 후위부는 전위부보다 +1되어있는 규칙 때문이다.

그러고, 현재 읽은 1인 비트를 0으로 취한다.

 

최종 누적값이 누적 합이다.

이를 [0, b]에 대해서 구하고, [0, a-1]에 대해서 구하면 된다.

 

코드

C++

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

ll solve(ll x, vector<ll> &pows) {
  ll res = x & 1; // 첫 비트는 별도로 구하기
  for (int i = 55; i > 0; --i) {
    if (x & ((ll)1 << i)) {
      res += pows[i - 1];
      res += (x - ((ll)1 << i) + 1);
      x &= ~((ll)1 << i);
    }
  }
  return res;
}

int main(int argc, char *argv[]) {
  ll a, b;
  cin >> a >> b;
  vector<ll> pows(55);
  pows[0] = 1;
  for (int i = 1; i < 55; ++i) {
    pows[i] = pows[i - 1] * 2 + ((ll)1 << i);
  }

  cout << solve(b, pows) - solve(a - 1, pows) << "\n";

  return 0;
}

/*
 * 0: 0000 -> 0
 * 1: 0001 -> 1
 * 2: 0010 -> 1
 * 3: 0011 -> 2
 * 4: 0100 -> 1
 * 5: 0101 -> 2
 * 6: 0110 -> 2
 * 7: 0111 -> 3
 * 8: 1000 -> 1
 * 9: 1001 -> 2
 * 10: 1010 -> 2
 * 11: 1011 -> 3
 * 12: 1100 -> 2
 * 13: 1101 -> 3
 * 14: 1110 -> 3
 * 15: 1111 -> 4