[백준 21758] 꿀 따기 C++

21758번 - 꿀 따기

문제 이해

정보올림피아드 2021 1차 대회 기출문제. 1차원 배열로 꿀의 양이 주어지고, 벌 두마리와 벌통을 배치한다. 벌은 벌통까지 직선으로 이동하고, 벌의 시작점에서는 꿀을 딸 수 없다.(다른 벌의 시작점도 포함)

접근

벌통과 꿀의 위치에 따라 세가지 경우로 나누어 볼 수 있다.

• 벌, 벌, 벌통 : 왼쪽 벌은 가장 왼쪽, 벌통은 가장 오른쪽에 고정해야 이득이다. 가운데 벌의 위치에 따라 따지 못하는 꿀이 달리지니, 가운데의 벌의 위치를 옮겨가며 탐색.
• 벌통, 벌, 벌 : 위와 마찬가지이다.
• 벌, 벌통, 벌 : 벌은 양 끝에 고정시킨 채, 벌통을 움직여가며 탐색.

탐색 과정을 살펴 보면 벌이 직선으로 이동하게 되니 구간의 합을 구해야 하는데, 누적합을 통해 시간을 줄여야 해결 가능하다.
누적합을 사용하지 않는다면, 한 경우를 탐색하는 데에 $O(N)$ 의 계산이 필요하고, 총 시간복잡도는 $O(N^2)$ 이 되는데, $N$ 은 최대 100,000이므로, $10^{10}$ 즉 100억번의 연산으로 100초 이상 소요된다.
누적합을 사용하게 되면 한 경우의 탐색에 $O(1)$, 총 시간복잡도는 $O(N)$으로 시간 내에 해결이 가능하다.

정답 코드

#include <bits/stdc++.h>
#define fastio             \
  ios::sync_with_stdio(0); \
  cin.tie(0);              \
  cout.tie(0)
#define MAX 100000

using namespace std;

int N;
int a[MAX], prefix[MAX];

void initPrefix() {
  prefix[0] = a[0];
  for (int i = 1; i < N; i++) {
    prefix[i] = prefix[i - 1] + a[i];
  }
}

int getSum(int i, int j) {
  if (i == 0)
    return prefix[j];
  return prefix[j] - prefix[i - 1];
}

void solve() {
  cin >> N;
  for (int i = 0; i < N; i++) {
    cin >> a[i];
  }
  initPrefix();

  int ans = 0;

  // 벌, 벌, 벌톨
  for (int i = 1; i < N - 1; i++) {
    ans = max(ans, getSum(1, N - 1) + getSum(i + 1, N - 1) - a[i]);
  }
  // 벌통, 벌, 벌
  for (int i = N - 2; i > 0; i--) {
    ans = max(ans, getSum(0, N - 2) - a[i] + getSum(0, i - 1));
  }
  // 벌, 벌통, 벌
  for (int i = 1; i < N - 1; i++) {
    ans = max(ans, getSum(1, i) + getSum(i, N - 2));
  }

  cout << ans;
}

int main() {
  fastio;

  int t = 1;
  // cin >> t;
  while (t--) {
    solve();
  }

  return 0;
}