제출 #1283645

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1283645		julsjlv	Tree (IOI24_tree)	C++20	18 / 100	2155 ms	1487488 KiB

tree

#include "tree.h"
#include <stdio.h>
#include <map>
#include <set>
#include <queue>
#define lol long long

using namespace std;

int n;
vector<int> p, w;
vector<vector<lol> > children;
map<lol, lol> counts;
lol zeroes;
// mode == 0: W <= 1
// mode == 1: W[P[i]] <= W[i] for each 1 <= i <= N
lol mode;
vector<int> initialGrades;

// return countLeaves
lol dfs(lol node)
{
  lol leaves = 0;
  if (w[node] == 0) {
    for (auto j: children[node]) {
      lol result = dfs(j);
      counts[result]++;
    }
    return 1;
  }
  for (auto j: children[node])
    leaves += dfs(j);
  return (children[node].size() == 0) ? 1 : leaves;
}

void init(vector<int> P, vector<int> W) {
  p = P;
  w = W;
  n = (int)p.size();
  counts.clear();
  children.clear();
  mode = 0;
  for (auto j: W) {
    if (j != 0 && j != 1) {
      mode = 1;
      break;
    }
  }
  children = vector<vector<lol> >(p.size(), vector<lol>());
  for (int i = 0; i < n; i++) {
    if (p[i] != -1)
      children[p[i]].push_back(i);
    zeroes += w[i] == 0;
  }
  
  lol leavesOfRoot = dfs(0);
  counts[leavesOfRoot]++;
  initialGrades = vector<int>(p.size(), 0);
  for (int i = 0; i < n; i++) {
    if (p[i] == -1)
      continue;
    initialGrades[p[i]]++;
  }
}

long long subtree(lol L, lol R, lol leavesCount) {
  
  lol sumLeaves = leavesCount * ((lol) L);
  if (sumLeaves <= R) {
    return sumLeaves;
  } else {
    return sumLeaves + (sumLeaves - R);
  }
}

long long query(int L, int R) {
  long long result = 0;
  if (mode == 0) {
    for (auto j: counts)
      result += subtree(L, R, j.first) * j.second;
    return result - zeroes * L;
  } else if (mode == 1) {
    vector<int> currentGrades = vector<int>(n, 0);
    vector<int> currentValues = vector<int>(n, 0);
    // cost, remaining
    vector<multiset<pair<int, int> > *> costsPerNode = vector<multiset<pair<int, int> > *>(n, NULL);
    queue<int> q;
    for (int i = 0; i < n; i++) {
      costsPerNode[i] = new multiset<pair<int, int> >();
      currentGrades[i] = initialGrades[i];
      if (currentGrades[i] == 0)
        q.push(i);
    }
    while (!q.empty()) {
      
      int i = q.front();
      q.pop();
      if (currentValues[i] < L) {
        result += w[i] * (L - currentValues[i]);
        currentValues[i] = L;
        if (R != L)
          costsPerNode[i]->insert({w[i], R - L});
      } else if (currentValues[i] <= R) {
        if (R != currentValues[i])
          costsPerNode[i]->insert({w[i], R - currentValues[i]});
      } else {
        while (currentValues[i] > R && !costsPerNode[i]->empty()) {
          pair<int, int> firstt = *costsPerNode[i]->begin();
          int toDecrease = min(currentValues[i] - R, firstt.second);
          currentValues[i] -= toDecrease;
          result += firstt.first * toDecrease;
          costsPerNode[i]->erase(costsPerNode[i]->begin());
          if (toDecrease != firstt.second)
            costsPerNode[i]->insert({firstt.first, firstt.second - toDecrease});
        }
        if (currentValues[i] > R) {
          result += w[i] * (currentValues[i] - R);
          currentValues[i] = R;
        }
      }

      if (i == 0)
        continue;
      
      currentGrades[p[i]]--;
      if (currentGrades[p[i]] == 0)
        q.push(p[i]);

      // unite costsPerNode[p[i]] with costsPerNode[i]
      if (costsPerNode[p[i]]->size() > costsPerNode[i]->size()) {
        for (auto j: (*costsPerNode[i]))
          costsPerNode[p[i]]->insert(j);
        costsPerNode[i] = costsPerNode[p[i]];
      } else {
        for (auto j: (*costsPerNode[p[i]]))
          costsPerNode[i]->insert(j);
        costsPerNode[p[i]] = costsPerNode[i];
      }
      currentValues[p[i]] += currentValues[i];
    }
    costsPerNode.clear();
    return result;
  }
  return -1;
}

• Subtask 0: 0 / 0
• Subtask 1: 0 / 10
• Subtask 2: 0 / 13
• Subtask 3: 0 / 18
• Subtask 4: 7 / 7
• Subtask 5: 11 / 11
• Subtask 6: 0 / 22
• Subtask 7: 0 / 19