recover.rsk.cpp

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

/* Recovers k increasing disjoint subsequences that cover the maximum possible
 * number of elements from A.  Runs in O(MN) time where M is the resulting
 * number of elements on all subsequences.  Based on the method
 * described in section 4 of Greene's "An Extension of Schensted's Theorem".
 * 
 * Expects A to be a permutation of [0, N - 1].
 */
vector<vector<int> > rsk_recover(const vector<int>& A, int k) {
  vector<vector<int>> h(k);
  int N = A.size();
  vector<tuple<int, int, int>> swaps;

  /* Run the RSK algorithm.  Record all of the swaps.  We only need to do
   * Type 3 swaps (which are Type 1 swaps when going back from canonical form).
   */
  for (int i = 0; i < N; i++) {
    int x = A[i];
    for (int j = 0; j < k; j++) {
      if (h[j].empty() || h[j].back() < x) {
        h[j].push_back(x);
        break;
      }

      int y;
      for (y = h[j].size() - 1; ; y--) {
        if (y == 0 || h[j][y - 1] < x) {
          swap(x, h[j][y]);
          break;
        }
        swaps.push_back(make_tuple(x, h[j][y - 1], h[j][y]));
      }
    }
  }

  /* Construct a linked list representing the k increasing subsequences
   * initially representing the canonical representation of A.
   */
  vector<int> nxt(N + 1, -1);
  vector<int> prv(N + 1, -1);
  for (int i = 0; i < k; i++) {
    prv[h[i][0]] = N;
    for (int j = 1; j < h[i].size(); j++) {
      prv[h[i][j]] = h[i][j - 1];
      nxt[h[i][j - 1]] = h[i][j];
    }
    nxt[h[i].back()] = N;
  }

  /* Replay the swaps backwards and adjust the subsequences appropriately. */
  for (int i = swaps.size() - 1; i >= 0; i--) {
    int x, y, z;
    tie(x, y, z) = swaps[i];
    if (nxt[x] != z) {
      /* Don't need to do anything if x and z aren't in the same subsequence. */
      continue;
    } else if (nxt[y] == -1) {
      /* Put y in place of x in its subsequence. */
      prv[y] = prv[x];
      nxt[prv[y]] = y;
      nxt[y] = z;
      prv[z] = y;
      prv[x] = nxt[x] = -1;
    } else {
      /* Splice lists; a->y->b and c->x->z->d becomes a->y->z->d and c->x->b. */
      nxt[x] = nxt[y];
      prv[nxt[x]] = x;
      nxt[y] = z;
      prv[z] = y;
    }
  }

  /* Reconstruct the actual subsequences from the linked list. */
  int cnt = 0;
  vector<int> seq(N, -1);
  vector<vector<int> > res(k);
  for (int i = 0; i < N; i++) {
    if (prv[i] != -1) {
      seq[i] = prv[i] == N ? cnt++ : seq[prv[i]];
      res[seq[i]].push_back(i);
    }
  }

  return res;
}

int main() {
  int N, k;
  cin >> N >> k;
  vector<int> A;
  for(int i = 0; i < N; i++) {
    int x;
    cin >> x;
    A.push_back(x);
  }

  vector<vector<int> > B = rsk_recover(A, k);
  for (int i = 0; i < B.size(); i++) {
    for (int j = 0; j < B[i].size(); j++) {
      if (j) cout << ' ';
      cout << B[i][j];
    }
    cout << endl;
  }

  return 0;
}