Submission #989282

#	Time	Username	Problem	Language	Result	Execution time	Memory
989282		OAleksa	Split the sequence (APIO14_sequence)	C++14	100 / 100	674 ms	111440 KiB

sequence

#include <bits/stdc++.h>
#define f first
#define s second
using namespace std;
const int N = 1e5 + 69;
const int K = 220;
int n, a[N], k, par[N][K];
long long dp[N], p[N];
struct line { 
  long long n, k;
  long long val(long long x) {
    return k * x + n;
  }
};
deque<pair<line, int>> dq[K];
signed main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
  int tt = 1;
  //cin >> tt;
  while (tt--) {
    cin >> n >> k;
    for (int i = 1;i <= n;i++)
      cin >> a[i];
    for (int i = 1;i <= n;i++)
      p[i] = p[i - 1] + a[i];
    auto Get = [&](int l, int r) {
      return p[r] - p[l - 1];
    };
    if (n <= 200) {
      vector<vector<long long>> dd(n + 1, vector<long long>(k + 1));
      for (int i = 0;i <= n;i++) {
        for (int j = 0;j <= k;j++)
          dd[i][j] = -1e18;
      }
      dq[0].push_back({{0, 0}, 0});
      dd[0][0] = 0;
      for (int i = 1;i < n;i++) {
        for (int j = 1;j <= min(i, k);j++) {  
          int x = p[n] - p[i], m = dq[j - 1].size();
          while (m > 1 && dq[j - 1][m - 1].f.val(x) <= dq[j - 1][m - 2].f.val(x)) {
            dq[j - 1].pop_back();
            --m;
          }
          dd[i][j] = dq[j - 1][m - 1].f.val(x) + p[i] * (p[n] - p[i]);
          line t = {dd[i][j], -p[i]};
          m = dq[j].size();
          while (m > 1 && (dq[j][0].f.k - t.k) * (dq[j][1].f.n - dq[j][0].f.n) >= (t.n - dq[j][0].f.n) * (dq[j][0].f.k - dq[j][1].f.k)) {
            dq[j].pop_front();
            --m;
          }
          dq[j].push_front({t, i});
        }
      }
      long long ans = 0;
      for (int i = 1;i < n;i++)
        ans = max(ans, dd[i][k]);
      vector<int> res;
      int t = -1;
      for (int i = 1;i < n;i++) {
        if (ans == dd[i][k])
          t = i;
      }
      assert(t != -1);
      res.push_back(t);
      for (int j = t - 1;j >= 1 && k > 0;j--) {
        long long f = dd[j][k - 1] + Get(j + 1, t) * Get(t + 1, n);
        if (f == dd[t][k]) {
          res.push_back(j);
          --k;
          t = j;
        }
      }
      cout << ans << '\n';
      reverse(res.begin(), res.end());
      for (auto j : res) 
        cout << j << ' ';
      return 0;
    }
    for (int i = 1;i <= n;i++) 
      dp[i] = -1e18;
    dq[0].push_back({{0, 0}, 0});
    for (int j = 1;j <= k;j++) {
      for (int i = 1;i < n;i++) {
        if (j > i)
          continue;
        int x = p[n] - p[i], m = dq[j - 1].size();
        while (m > 1 && dq[j - 1][m - 1].f.val(x) <= dq[j - 1][m - 2].f.val(x)) {
          dq[j - 1].pop_back();
          --m;
        }
        dp[i] = dq[j - 1][m - 1].f.val(x) + p[i] * (p[n] - p[i]);
        par[i][j] = dq[j - 1][m - 1].s;
        line t = {dp[i], -p[i]};
        m = dq[j].size();
        while (m > 1 && (dq[j][0].f.k - t.k) * (dq[j][1].f.n - dq[j][0].f.n) >= (t.n - dq[j][0].f.n) * (dq[j][0].f.k - dq[j][1].f.k)) {
          dq[j].pop_front();
          --m;
        }
        dq[j].push_front({t, i});
      }
    }
    long long ans = 0;
    for (int i = 1;i < n;i++)
      ans = max(ans, dp[i]);
    vector<int> res;
    int t = -1;
    for (int i = 1;i < n;i++) {
      if (ans == dp[i])
        t = i;
    }
    assert(t != -1);
    int j = k;
    while (t != 0) {
      res.push_back(t);
      t = par[t][j--];
    }
    cout << ans << '\n';
    assert(res.size() == k);
    reverse(res.begin(), res.end());
    for (auto j : res) 
      cout << j << ' ';
  }
  return 0; 
}

Compilation message (stderr)

In file included from /usr/include/c++/10/cassert:44,
                 from /usr/include/x86_64-linux-gnu/c++/10/bits/stdc++.h:33,
                 from sequence.cpp:1:
sequence.cpp: In function 'int main()':
sequence.cpp:120:23: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  120 |     assert(res.size() == k);
      |            ~~~~~~~~~~~^~~~

• Subtask 1: 11 / 11
• Subtask 2: 11 / 11
• Subtask 3: 11 / 11
• Subtask 4: 17 / 17
• Subtask 5: 21 / 21
• Subtask 6: 29 / 29