[엘리스 코드 챌린지 예선 Day 6] 빨간 선과 파란 선

문제 설명[링크]

제한 시간: 1 초

N개의 정점이 있다.
차례마다 다음을 반복해서 N개의 정점 사이에 간선을 연결하려고 한다.

먼저 2개의 서로 연결되지 않은 정점 u와 v를 고른다.
그 이후, u가 포함된 연결 요소의 모든 정점들 각각에서, v가 포함된 연결 요소의 모든 정점들 각각으로 간선을 추가한다.
간선의 색은 빨간색 혹은 파란색이다.
k번째 차례에 사용할 색깔이 주어질 때, 정점을 골라서 얻을 수 있는 빨간 간선 개수의 최솟값을 구하여라.

입력
첫 번째 줄에 정점의 수 N이 주어진다.
2≤N≤30
두 번째 줄에 각 차례에 사용할 색깔을 표시하는 N−1개의 수가 공백을 구분으로 주어진다.
숫자가 0이면 빨간 간선을, 1이면 파란 간선을 사용한다는 뜻이다.
입력되는 모든 수들은 정수이다.
출력
문제의 조건을 만족하면서 간선을 연결할 때, 얻을 수 있는 빨간 간선 개수의 최솟값을 출력한다.

입력 예시
5
1 1 0 1

출력 예시
1

풀이

• 맨 마지막 상태는 결국 n개의 정점 모두 연결된 상태 {30}
• 맨 마지막 부터 순회하여, 연결 그룹 하나를 잡고, 2개로 나누면서 진행
• 맨 마지막 연결이 파란색(1) 이라면, 힙에 1푸쉬
• 아니라면, 힙의 최소값과 남은 n값을 비교하여, 추가되는 연결 간선 수의 최솟값을 비교하여 연산 진행
  - 만약 힙에서 빠졌다면, 해당값 + 1을 힙에 다시 저장
  
• O(nlogn)

#include <iostream>
#include <queue>
#include <vector>
using namespace std;

#define FASTIO                 \
  ios::sync_with_stdio(false); \
  cin.tie(NULL);               \
  cout.tie(NULL);

#define F(i, a, b) for (int i = a; i < b; i += 1)
#define FR(i, b, a) for (int i = b; i >= a; i -= 1)

int n;
int result = 0;
vector<int> arr;
priority_queue<int, vector<int>, greater<int>> minHeap;

int main() {
  FASTIO
  cin >> n;
  arr.resize(n - 1);
  F(i, 0, n - 1) cin >> arr[i];

  FR(i, n - 2, 0) {
    if (arr[i] == 1) {
      minHeap.push(1);
    } else {
      const int compare1 = minHeap.size() ? minHeap.top() : 101010101;
      const int compare2 = n - 1;

      if (compare1 <= compare2) {
        result += compare1;
        minHeap.pop();
        minHeap.push(compare1 + 1);
      } else {
        result += compare2;
      }
    }

    n -= 1;
  }

  cout << result;
  return 0;
}