Submission #1196267

#TimeUsernameProblemLanguageResultExecution timeMemory
1196267avighnaHoliday (IOI14_holiday)C++20
100 / 100
3820 ms16020 KiB
#include <algorithm>
#include <functional>
#include <vector>

template <typename T> class SegmentTree {
public:
  int n;
  T idt;
  std::vector<T> seg;
  std::function<T(T, T)> f;

  SegmentTree(int n, std::function<T(T, T)> f = std::plus<T>(), T idt = T()) : n(n), idt(idt), f(f), seg(2 * n, idt) {}

  void set(int idx, T x) {
    for (seg[idx += n] = x, idx /= 2; idx > 0; idx /= 2) {
      seg[idx] = f(seg[2 * idx], seg[2 * idx + 1]);
    }
  }

  T query(int l, int r) {
    T ansL = idt, ansR = idt;
    for (l += n, r += n + 1; l < r; l /= 2, r /= 2) {
      if (l & 1) ansL = f(ansL, seg[l++]);
      if (r & 1) ansR = f(seg[--r], ansR);
    }
    return f(ansL, ansR);
  }

  int partition_point(int l, const std::function<bool(T)> &t) {
    T p = idt;
    for (l += n; t(f(p, seg[l])) and l; l /= 2) {
      if (l & 1 and l != 1) p = f(p, seg[l++]);
    }
    if (!l) {
      return n;
    }
    while (l < n) {
      if (t(f(p, seg[l <<= 1]))) p = f(p, seg[l++]);
    }
    return l - n;
  }
};


long long int solve(int n, int start, int d, std::vector<int> attraction) {
  auto build_a = [&]() {
    std::vector<std::pair<int, int>> a;
    for (int i = 0; i < n; ++i) {
      a.push_back({attraction[i], i});
    }
    std::sort(a.rbegin(), a.rend());
    std::vector<std::pair<int, int>> new_a(n);
    for (int i = 0; i < n; ++i) {
      new_a[a[i].second] = {a[i].first, i};
    }
    return a = std::move(new_a);
  };
  std::vector<std::pair<int, int>> a = build_a();
  struct Node {
    int count = 0;
    int64_t sum = 0;
  };
  auto round_up_pow2 = [](int n) {
    --n;
    n |= n >> 1;
    n |= n >> 2;
    n |= n >> 4;
    n |= n >> 8;
    n |= n >> 16;
    return n + 1;
  };  
  SegmentTree<Node> st(round_up_pow2(n), [&](Node a, Node b) {
    return Node({a.count + b.count, a.sum + b.sum});
  });

  auto calc_f = [&](auto &&self, int tl, int tr, int l, int r, int start, std::vector<std::pair<int, int64_t>> &f) {
    int tm = (tl + tr) / 2;
    std::pair<int64_t, int> best = {-1, -1};
    for (int i = l; i <= r; ++i) {
      st.set(a[i].second, {1, a[i].first});
      int idx = st.partition_point(0, [&](Node x) {
        return x.count <= tm - (i - start);
      }) - 1;
      if (idx == -1) {
        continue;
      }
      best = std::max(best, {st.query(0, idx).sum, -i});
    }
    best.second = -best.second;
    f[tm] = {best.second, best.first};
    if (tl == tr) {
      return;
    }
    for (int i = f[tm].first; i <= r; ++i) {
      st.set(a[i].second, Node{});
    }
    self(self, tm + 1, tr, f[tm].first, r, start, f);
    for (int i = l; i <= r; ++i) {
      st.set(a[i].second, Node{});
    }
    self(self, tl, tm, l, f[tm].first, start, f);
  };

  std::vector<std::pair<int, int64_t>> f1(d + 1);
  calc_f(calc_f, 1, d, start, n - 1, start, f1);
  for (int i = 0; i < n; ++i) {
    st.set(i, Node{});
  }
  std::reverse(attraction.begin(), attraction.begin() + start);
  for (int i = start; i < n; ++i) {
    attraction[i] = 0;
  }
  a = build_a();
  std::vector<std::pair<int, int64_t>> f2(d + 1);
  calc_f(calc_f, 1, d, 0, n - 1, 0, f2);
  long long int ans = 0;
  for (int i = 1; i <= d; ++i) {
    long long int cur = f1[i].second;
    int rem = d - i - f1[i].first + start - 1;
    if (rem > 0) {
      cur += f2[rem].second;
    }
    ans = std::max(ans, cur);
  }
  return ans;
}

long long int findMaxAttraction(int n, int start, int d, int attraction[]) {
  std::vector<int> a(n);
  for (int i = 0; i < n; ++i) {
    a[i] = attraction[i];
  }
  long long int ans1 = solve(n, start, d, a);
  std::reverse(a.begin(), a.end());
  start = n - start - 1;
  long long int ans2 = solve(n, start, d, a);
  return std::max(ans1, ans2);
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...