Submission #1105032

#	Time	Username	Problem	Language	Result	Execution time	Memory
1105032		flashmt	Tree (IOI24_tree)	C++17	41 / 100	2119 ms	853592 KiB

tree

// 41 points
#include "tree.h"
#include <bits/stdc++.h>
using namespace std;
const int N = 2e5 + 5;
const long long oo = 1LL << 40;

int n, leaf, isSub45;
vector<int> p, w, a[N], subtree, roots;
map<long long, long long> m[N];
vector<long long> sum, prefix;

void init(int x)
{
  if (empty(a[x]))
  {
    subtree[x] = 1;
    leaf += w[x];
  }
  else
  {
    for (int y : a[x])
      init(y);

    if (!w[x])
    {
      subtree[x] = 1;
      for (int y : a[x])
        if (subtree[y] > 1)
          roots.push_back(subtree[y]);
    }
    else
    {
      for (int y : a[x])
        subtree[x] += subtree[y];
      if (x == 0)
        roots.push_back(subtree[0]);
    }
  }
}

void init(vector<int> P, vector<int> W) {
  p = P;
  w = W;
  isSub45 = 1;
  for (int x : w)
    if (x > 1)
      isSub45 = 0;
  n = (int)p.size();
  for (int i = 0; i < n; i++)
    a[i].clear();
  for (int i = 1; i < n; i++)
    a[p[i]].push_back(i);
  leaf = 0;
  subtree = vector<int>(n);
  sum = vector<long long>(n, 0);
  roots.clear();
  init(0);
  sort(begin(roots), end(roots));
  prefix.clear();
  prefix.push_back(roots[0]);
  for (int i = 1; i < size(roots); i++)
    prefix.push_back(prefix[i - 1] + roots[i]);
}

long long calc(int x, int low, int high)
{
  m[x].clear();
  if (empty(a[x]))
  {
    sum[x] = low;
    return 1LL * w[x] * low;
  }

  long long res = 0;
  sum[x] = 0;
  m[x][w[x]] = oo;
  for (int y : a[x])
  {
    res += calc(y, low, high);
    sum[x] += sum[y];
    for (auto [k, v] : m[y])
      m[x][k] += v;
  }

  if (sum[x] > high)
  {
    long long need = sum[x] - high;
    for (auto [curW, quota] : m[x])
    {
      long long use = min(need, quota);
      res += use * curW;
      need -= use;
      m[x][curW] -= use;
      if (!need)
        break;
    }
    assert(need == 0);
    sum[x] = high;
  }

  map<long long, long long> newM;
  long long cur = sum[x] - low;
  for (auto [curW, quota] : m[x])
  {
    long long use = min(quota, cur);
    newM[curW] += use;
    cur -= use;
    if (!cur)
      break;
  }
  swap(m[x], newM);

  return res;
}

long long query(int low, int high) {
  if (isSub45)
  {
    long long ans = 1LL * leaf * low;
    int sz = size(roots);
    int l = 0, r = sz - 1, res = sz;
    while (l <= r)
    {
      int m = (l + r) / 2;
      if (1LL * roots[m] * low > high)
      {
        res = m;
        r = m - 1;
      }
      else l = m + 1;
    }
    if (res < sz)
    {
      ans += 1LL * (prefix[sz - 1] - (res ? prefix[res - 1] : 0)) * low;
      ans -= 1LL * (sz - res) * high;
    }
    return ans;
  }
  return calc(0, low, high);
}

Compilation message (stderr)

tree.cpp: In function 'void init(std::vector<int>, std::vector<int>)':
tree.cpp:62:21: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   62 |   for (int i = 1; i < size(roots); i++)
      |                   ~~^~~~~~~~~~~~~

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