Submission #3665432

---

Source Code Expand

#include<bits/stdc++.h>

using namespace std;

using int64 = long long;
const int64 INF = 1LL << 40;

struct UnionFind {
  vector< int > data;

  UnionFind(int sz) {
    data.assign(sz, -1);
  }

  bool unite(int x, int y) {
    x = find(x), y = find(y);
    if(x == y) return (false);
    if(data[x] > data[y]) swap(x, y);
    data[x] += data[y];
    data[y] = x;
    return (true);
  }

  int find(int k) {
    if(data[k] < 0) return (k);
    return (data[k] = find(data[k]));
  }

  int size(int k) {
    return (-data[find(k)]);
  }
};

int64 dp5[5001][5001][2][2];
int D[5000];
bool v[5001];


int main() {
  int K, N, A[5000], B[5000];
  cin >> N >> K;
  for(int i = 0; i < N; i++) cin >> A[i];
  for(int i = 0; i < N; i++) cin >> B[i];
  vector< int > ord(N);
  iota(begin(ord), end(ord), 0);
  sort(begin(ord), end(ord), [&](int a, int b) {
    return A[a] < A[b];
  });

  memset(dp5, -1, sizeof(dp5));
  UnionFind uf(N);
  for(int i = 0; i < N; i++) {
    D[i] = B[ord[i]];
    --D[i];
    uf.unite(D[i], i);
  }


  int sz = 0;
  vector< int64 > dp(sz + 1, -INF);
  dp[0] = 0;
  for(int i = 0; i < N; i++) {
    if(v[i]) continue;
    v[i] = true;
    dp5[D[i]][1][0][1] = i + 1;
    dp5[D[i]][1][1][0] = D[i] + 1;
    int idx = D[i], szsz = 1;
    while(idx != i) {
      v[idx] = true;
      int nxt = D[idx];
      for(int j = 0; j <= szsz; j++) {
        for(int k = 0; k < 2; k++) {
          for(int l = 0; l < 2; l++) {
            if(dp5[idx][j][k][l] == -1) continue;

            if(nxt == i) { // 最後の頂点
              if(l) { // 最後の頂点はカウント済み
                if(k) { // この頂点はカウント済み
                  dp5[nxt][j][true][l] = max(dp5[nxt][j][true][l], dp5[idx][j][k][l] + (nxt + 1));
                  dp5[nxt][j][false][l] = max(dp5[nxt][j][false][l], dp5[idx][j][k][l] + (idx + 1));
                } else { // この頂点は未カウント
                  dp5[nxt][j][true][l] = max(dp5[nxt][j][true][l], dp5[idx][j][k][l] + (nxt + 1));
                  dp5[nxt][j + 1][false][l] = max(dp5[nxt][j + 1][false][l], dp5[idx][j][k][l] + (idx + 1));
                }
              } else { // 最後の頂点は未カウント
                if(k) { //ここに入る この頂点はカウント済み
                  dp5[nxt][j + 1][true][l] = max(dp5[nxt][j + 1][true][l], dp5[idx][j][k][l] + (nxt + 1));
                  dp5[nxt][j][false][l] = max(dp5[nxt][j][false][l], dp5[idx][j][k][l] + (idx + 1));
                } else {
                  dp5[nxt][j + 1][true][l] = max(dp5[nxt][j + 1][true][l], dp5[idx][j][k][l] + (nxt + 1));
                  dp5[nxt][j + 1][false][l] = max(dp5[nxt][j + 1][false][l], dp5[idx][j][k][l] + (idx + 1));
                }
              }
            } else if(k) { // この頂点は既にカウント済み
              dp5[nxt][j + 1][true][l] = max(dp5[nxt][j + 1][true][l], dp5[idx][j][k][l] + (nxt + 1));
              dp5[nxt][j][false][l] = max(dp5[nxt][j][false][l], dp5[idx][j][k][l] + (idx + 1));
            } else { // この頂点は未カウント
              dp5[nxt][j + 1][true][l] = max(dp5[nxt][j + 1][true][l], dp5[idx][j][k][l] + (nxt + 1));
              dp5[nxt][j + 1][false][l] = max(dp5[nxt][j + 1][false][l], dp5[idx][j][k][l] + (idx + 1));
            }
          }
        }
      }
      idx = D[idx];
      ++szsz;
    }

    vector< int64 > dp2(uf.size(i) + 1, -INF);
    for(int j = 0; j <= uf.size(i); j++) {
      for(int k = 0; k < 2; k++) {
        for(int l = 0; l < 2; l++) {
          if(dp5[i][j][k][l] == -1) continue;
          dp2[j] = max(dp2[j], dp5[i][j][k][l]);
        }
      }
    }
    vector< int64 > dp3(sz + uf.size(i) + 1, -INF);
    for(int j = 0; j <= sz; j++) {
      for(int k = uf.size(i); k >= 0; k--) {
        dp3[j + k] = max(dp3[j + k], dp[j] + dp2[k]);
      }
    }
    sz += uf.size(i);
    // うくちゃんか？
    dp.swap(dp3);
  }

  int64 uku = 0;
  for(int i = K; i < dp.size(); i++) {
    uku = max(uku, dp[i]);
  }
  cout << uku << endl;
}

Submission Info

Judge Result