제출 #1196264

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1196264		avighna	Holiday (IOI14_holiday)	C++20	100 / 100	3832 ms	16020 KiB

holiday

#include <vector>
#include <cassert>
#include <functional>
#include <algorithm>

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) {
  // create a segtree to store the elements
  // you'll have to sort them and assign an index
  // then walk down the tree to find the sum of the first k largest elements
  // first sort
  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();
  // segtree
  struct Node {
    int count = 0;
    int64_t sum = 0;
  };
  auto round_up_pow2 = [](int n) {
    if (n == 0) return 1;
    --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});
  });

  // for [l, r], youll need to iteratively add everything to segtree. then for [m,r] you first remove the right values (and repeat).
  // for [l, m] remove the left values and newly set right values (and repeat)

  // given starting s, define f(x) with domain [1, n] as the city to jump to (given x stamina) so that the sum you get is optimal
  // property: f(x) <= f(x + 1). any city to the left of and including f(x) will be more optimal with x+1 stamina, so f(x+1)
  //     could either be f(x) or it could be some newer city down the line that improves it
  // now we calculate f(x) [1, n] in O(n log^2 n) time with dnc
  // to compute for [tl, tr] first compute for tm
  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;
    // comute f(tm)
    std::pair<int64_t, int> best = {-1, -1};
    for (int i = l; i <= r; ++i) {
      st.set(a[i].second, {1, a[i].first});
      // tm-(i-st) stamina: so pick (at max) tm-(i-st) largest vals
      int idx = st.partition_point(0, [&](Node x) {
        return x.count <= tm - (i - start);
      }) - 1;
      // last index with count <= tm - (i-st)
      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;
    }
    // remove right values to prepare for recursive call
    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);
    // remove all values for left recursive call
    for (int i = l; i <= r; ++i) {
      st.set(a[i].second, Node{});
    }
    self(self, tl, tm, l, f[tm].first, start, f);
  };

  // go from start to right then back to start-1 and left from there
  // so first calculate f1(x) ==> starting point going right
  std::vector<std::pair<int, int64_t>> f1(d + 1);
  calc_f(calc_f, 1, d, start, n - 1, start, f1);
  // reset segtree
  for (int i = 0; i < n; ++i) {
    st.set(i, Node{});
  }
  // then calculate f2(x) ==> starting point (start-1) going left. to do this just reverse the array
  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);
  // now iterate over all possible stamina values dedicated to right from 1 to d
  long long int ans = 0;
  for (int i = 1; i <= d; ++i) {
    // we will give i stamina to right and the remaining d-i-[whatever stamina was used to move to start-1] to left
    // if it is nonzero ofc
    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);
}

• Subtask 1: 7 / 7
• Subtask 2: 23 / 23
• Subtask 3: 17 / 17
• Subtask 4: 53 / 53