BOJ 17501. 수식 트리

솔직히 왜 골드2인지 모르겠는 문제. 간단한 트리 DP?를 통해 해결했다.

문제

solved.ac: Gold II (2024.09.06)

 

어차피 식을 모두 합치면, (a1 + a2 + ... ) - (b1 + b2 + ... ) 꼴이 될 것이다.

따라서 최대한 작은 수들을 빼면 된다. 이건 정렬 + 그리디로 된다.

 

어떤 서브트리와 두 서브트리에서 빼야 하는 노드의 개수를 각각 A, B개라 하자.

만약 현재 노드가 +라면, 이 서브트리에서 빼야 하는 노드의 개수는 자명히 A + B개가 된다.

 

반면 현재 노드가 -라면, 원래 빼야 했던 노드는 다 더해지고, 원래 더해야 했던 노드들이 빼지게 된다.

따라서 우측 서브트리의 트리 크기를 sz라 하면, 빼야 하는 노드의 개수는 A + (sz - B)개가 된다.

 

그러면 총 빼야 하는 노드의 수 s를 구할 수 있고, 오름차순 정렬 후 앞에서부터 s개는 빼고, 나머지는 더하면 된다.

#include <iostream>
#include <vector>
#include <string>
#include <algorithm>

using namespace std;

int n;
int dp[200001], sz[200001];
string tree[200001];

int child[200001][2];

void DP(int cur)
{
  if(cur <= n)
  {
    sz[cur] = 1;
    dp[cur] = 0;
    return;
  }

  int l = child[cur][0], r = child[cur][1];
  DP(l); DP(r);

  sz[cur] = sz[l] + sz[r];
  if(tree[cur] == "+") dp[cur] = dp[l] + dp[r];
  else dp[cur] = dp[l] + (sz[r] - dp[r]);
}

int main()
{
  ios::sync_with_stdio(false);
  cin.tie(0);

  cin >> n;

  vector<int> arr(n);
  for(int i = 0; i < n; i++) cin >> arr[i];
  for(int i = n + 1; i < 2 * n; i++) cin >> tree[i] >> child[i][0] >> child[i][1];

  DP(2 * n - 1);
  sort(arr.begin(), arr.end());

  int res = 0, idx = dp[2 * n - 1];
  for(int i = 0; i < n; i++)
    {
      if(i < idx) res -= arr[i];
      else res += arr[i];
    }
  cout << res;
}