Problem p03021 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef pair<int, int> ii;

const int N = 2000 + 5;

int n;
char s[N];
vector<int> v[N];

int sum;

void dfs(int p, int x, int d) {
  if (s[x] == '1') {
    sum += d;
  }
  for (auto u : v[x]) {
    if (u != p) {
      dfs(x, u, d + 1);
    }
  }
}

priority_queue<int> q[N];

void dfs2(int p, int x, int d, int &need) {
  if (s[x] == '1') {
    q[x].push(d);
  }
  for (auto u : v[x]) {
    if (u != p) {
      dfs2(x, u, d + 1, need);
      while (q[x].size() and q[x].top() > d and q[u].size() and
             q[u].top() > d) {
        int v = -1;

        v = q[x].top();
        q[x].pop();
        q[x].push(v - 1);

        v = q[u].top();
        q[u].pop();
        q[u].push(v - 1);

        need -= 2;
      }
      while (q[u].size()) {
        q[x].push(q[u].top());
        q[u].pop();
      }
    }
  }
}

/*
6
101011
1 2
2 3
2 4
4 5
4 6

7
1111000
1 7
2 7
6 7
5 6
3 5
4 5
*/

int main() {
  scanf("%d", &n);
  scanf("%s", s + 1);
  for (int i = 1; i < n; i++) {
    int x, y;
    scanf("%d %d", &x, &y);
    v[x].push_back(y);
    v[y].push_back(x);
  }
  int ans = 1e9;
  for (int i = 1; i <= n; i++) {
    // vector<priority_queue<int>> q(v[i].size());
    // priority_queue<ii> vals;
    priority_queue<int> vals;
    int id = 0;
    int sum_all = 0;
    vector<ii> sums;
    for (auto u : v[i]) {
      sum = 0;
      dfs(i, u, 1);
      sums.push_back({sum, u});
      sum_all += sum;
      // if(!q[id].empty()) {
      //   vals.push({q[id].top(), id});
      // }
      if (sum) {
        vals.push(sum);
      }
      id++;
    }
    // auto del = [&](int x) {
    //   int top = q[x].top();
    //   q[x].pop();
    //   if(top > 1)
    //     q[x].push(top - 1);
    //   // printf("pop %d->%d\n", x, top);
    //   if(q[x].size()) {
    //     // printf("add %d\n", q[x].top());
    //     vals.push({q[x].top(), x});
    //   }
    // };

    // int res = 0;
    // while(vals.size() > 1) {
    //   // int x = vals.top().second; vals.pop();
    //   // int y = vals.top().second; vals.pop();
    //   int x = vals.top(); vals.pop();
    //   int y = vals.top(); vals.pop();
    //   // del(x);
    //   // del(y);
    //   if(x > 1)
    //     vals.push(x - 1);
    //   if(y > 1)
    //     vals.push(y - 1);
    //   res++;
    // }
    // if(vals.empty())
    //   ans = min(ans, res);
    if (sum_all % 2)
      continue;
    id = -1;
    int id_sum = 0;
    for (auto x : sums) {
      if (x.first > sum_all - x.first) {
        id_sum = x.first;
        id = x.second;
      }
    }
    if (id == -1) {
      ans = min(ans, sum_all / 2);
      continue;
    }
    int need = id_sum - (sum_all - id_sum);
    for (int i = 1; i <= n; i++) {
      while (q[i].size())
        q[i].pop();
    }
    dfs2(i, id, 1, need);
    if (need <= 0)
      ans = min(ans, sum_all / 2);
  }
  printf("%d\n", ans > 8e8 ? -1 : ans);
  return 0;
}

Fixed submission:
#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef pair<int, int> ii;

const int N = 2000 + 5;

int n;
char s[N];
vector<int> v[N];

int sum;

void dfs(int p, int x, int d) {
  if (s[x] == '1') {
    sum += d;
  }
  for (auto u : v[x]) {
    if (u != p) {
      dfs(x, u, d + 1);
    }
  }
}

priority_queue<int> q[N];

void dfs2(int p, int x, int d, int &need) {
  if (s[x] == '1') {
    q[x].push(d);
  }
  for (auto u : v[x]) {
    if (u != p) {
      dfs2(x, u, d + 1, need);
      while (q[x].size() and q[x].top() > d and q[u].size() and
             q[u].top() > d) {
        int v = -1;

        v = q[x].top();
        q[x].pop();
        q[x].push(v - 1);

        v = q[u].top();
        q[u].pop();
        q[u].push(v - 1);

        need -= 2;
      }
      if (q[u].size() > q[x].size())
        swap(q[x], q[u]);
      while (q[u].size()) {
        q[x].push(q[u].top());
        q[u].pop();
      }
    }
  }
}

/*
6
101011
1 2
2 3
2 4
4 5
4 6

7
1111000
1 7
2 7
6 7
5 6
3 5
4 5
*/

int main() {
  scanf("%d", &n);
  scanf("%s", s + 1);
  for (int i = 1; i < n; i++) {
    int x, y;
    scanf("%d %d", &x, &y);
    v[x].push_back(y);
    v[y].push_back(x);
  }
  int ans = 1e9;
  for (int i = 1; i <= n; i++) {
    // vector<priority_queue<int>> q(v[i].size());
    // priority_queue<ii> vals;
    priority_queue<int> vals;
    int id = 0;
    int sum_all = 0;
    vector<ii> sums;
    for (auto u : v[i]) {
      sum = 0;
      dfs(i, u, 1);
      sums.push_back({sum, u});
      sum_all += sum;
      // if(!q[id].empty()) {
      //   vals.push({q[id].top(), id});
      // }
      if (sum) {
        vals.push(sum);
      }
      id++;
    }
    // auto del = [&](int x) {
    //   int top = q[x].top();
    //   q[x].pop();
    //   if(top > 1)
    //     q[x].push(top - 1);
    //   // printf("pop %d->%d\n", x, top);
    //   if(q[x].size()) {
    //     // printf("add %d\n", q[x].top());
    //     vals.push({q[x].top(), x});
    //   }
    // };

    // int res = 0;
    // while(vals.size() > 1) {
    //   // int x = vals.top().second; vals.pop();
    //   // int y = vals.top().second; vals.pop();
    //   int x = vals.top(); vals.pop();
    //   int y = vals.top(); vals.pop();
    //   // del(x);
    //   // del(y);
    //   if(x > 1)
    //     vals.push(x - 1);
    //   if(y > 1)
    //     vals.push(y - 1);
    //   res++;
    // }
    // if(vals.empty())
    //   ans = min(ans, res);
    if (sum_all % 2)
      continue;
    id = -1;
    int id_sum = 0;
    for (auto x : sums) {
      if (x.first > sum_all - x.first) {
        id_sum = x.first;
        id = x.second;
      }
    }
    if (id == -1) {
      ans = min(ans, sum_all / 2);
      continue;
    }
    int need = id_sum - (sum_all - id_sum);
    for (int i = 1; i <= n; i++) {
      while (q[i].size())
        q[i].pop();
    }
    dfs2(i, id, 1, need);
    if (need <= 0)
      ans = min(ans, sum_all / 2);
  }
  printf("%d\n", ans > 8e8 ? -1 : ans);
  return 0;
}