6.2.2 Enumerating Arrangements

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179

/*

For the purposes of this section, we define a "size $k$ arrangement of $n$" to be a permutation of a
size $k$ subset of the integers in $[0, n)$, for $0 \leq k \leq n$. There are
$n \mathbin{\text{permute}} k$ possible arrangements, but $n^k$ possible arrangements if repeated
values are allowed.

For arrangements without repeats, each prefix fixes some used values and leaves a shrinking set of
choices for the remaining positions. Ranking counts how many unused smaller choices could have been
placed at the current position, multiplied by the number of possible suffix arrangements; unranking
performs the inverse search by subtracting those block sizes. With repeats allowed, the sequence is
just a base-$n$ counter.

- `next_arrangement(n, a)` tries to rearrange `a` to the next lexicographically greater arrangement,
  returning true if such an arrangement exists or false if the array is already in descending order
  (in which case `a` is unchanged). The input `a` must consist of distinct integers in the range
  $[0, `n`)$.
- `arrangement_by_rank(n, k, r)` returns the size $k$ arrangement of $n$ which is lexicographically
  ranked $r$ out of all size $k$ arrangements of $n$, where $r$ is a 0-based rank in the range
  $[0, n \mathbin{\text{permute}} k)$.
- `rank_by_arrangement(n, a)` returns an integer representing the 0-based rank of arrangement `a`,
  which must consist of distinct integers in the range $[0, `n`)$.
- `next_arrangement_with_repeats(n, a)` tries to rearrange `a` to the next lexicographically greater
  arrangement with repeats, returning true if such an arrangement exists or false if the array is
  already in descending order (in which case `a` is unchanged). The input `a` must consist of
  integers in the range $[0, `n`)$. If `a` were interpreted as a $k$ digit integer in base $n$, this
  function could be thought of as incrementing the integer.

Time Complexity:
- O(n*k) per call to `next_arrangement()`, `arrangement_by_rank()`, and `rank_by_arrangement(n, a)`,
  where $k$ is the size of `a`.
- O(k) per call to `next_arrangement_with_repeats()`.

Space Complexity:
- O(n) auxiliary for `next_arrangement()`, `arrangement_by_rank()`, and `rank_by_arrangement()`.
- O(1) auxiliary for `next_arrangement_with_repeats()`.

*/

#include <algorithm>
#include <cstdint>
#include <numeric>
#include <vector>

bool next_arrangement(int n, std::vector<int> &a) {
  int k = static_cast<int>(a.size());
  std::vector<char> used(n);
  for (int i = 0; i < k; i++) {
    used[a[i]] = true;
  }
  for (int i = k - 1; i >= 0; i--) {
    used[a[i]] = false;
    for (int j = a[i] + 1; j < n; j++) {
      if (!used[j]) {
        a[i++] = j;
        used[j] = true;
        for (int x = 0; i < k; x++) {
          if (!used[x]) {
            a[i++] = x;
          }
        }
        return true;
      }
    }
  }
  return false;
}

int64_t n_permute_k(int n, int k) {
  int64_t res = 1;
  for (int i = 0; i < k; i++) {
    res *= n - i;
  }
  return res;
}

std::vector<int> arrangement_by_rank(int n, int k, int64_t r) {
  std::vector<int> values(n), res(k);
  std::iota(values.begin(), values.end(), 0);
  for (int i = 0; i < k; i++) {
    int64_t count = n_permute_k(n - 1 - i, k - 1 - i);
    int pos = r / count;
    res[i] = values[pos];
    values.erase(values.begin() + pos);
    r %= count;
  }
  return res;
}

int64_t rank_by_arrangement(int n, const std::vector<int> &a) {
  int k = static_cast<int>(a.size());
  int64_t res = 0;
  std::vector<char> used(n);
  for (int i = 0; i < k; i++) {
    int count = 0;
    for (int j = 0; j < a[i]; j++) {
      if (!used[j]) {
        count++;
      }
    }
    res += count * n_permute_k(n - i - 1, k - i - 1);
    used[a[i]] = true;
  }
  return res;
}

bool next_arrangement_with_repeats(int n, std::vector<int> &a) {
  int k = static_cast<int>(a.size());
  for (int i = k - 1; i >= 0; i--) {
    if (a[i] < n - 1) {
      a[i]++;
      std::fill(a.begin() + i + 1, a.end(), 0);
      return true;
    }
  }
  return false;
}

/*** Example Usage and Output:

4 permute 3 arrangements:
{0,1,2} {0,1,3} {0,2,1} {0,2,3} {0,3,1} {0,3,2} {1,0,2} {1,0,3} {1,2,0} {1,2,3}
{1,3,0} {1,3,2} {2,0,1} {2,0,3} {2,1,0} {2,1,3} {2,3,0} {2,3,1} {3,0,1} {3,0,2}
{3,1,0} {3,1,2} {3,2,0} {3,2,1}

4^2 arrangements with repeats:
{0,0} {0,1} {0,2} {0,3} {1,0} {1,1} {1,2} {1,3} {2,0} {2,1} {2,2} {2,3} {3,0}
{3,1} {3,2} {3,3}

***/

#include <cassert>
#include <iostream>
using namespace std;

template<typename It>
void print_range(It lo, It hi) {
  cout << "{";
  for (; lo != hi; ++lo) {
    cout << *lo << (lo == hi - 1 ? "" : ",");
  }
  cout << "} ";
}

int main() {
  {
    int n = 4, k = 3;
    vector<int> a{0, 1, 2};
    cout << n << " permute " << k << " arrangements:" << endl;
    int count = 0;
    do {
      print_range(a.begin(), a.end());
      assert(a == arrangement_by_rank(n, k, count));
      assert(rank_by_arrangement(n, a) == count);
      count++;
      if (count == 10 || count == 20) {
        cout << endl;
      }
    } while (next_arrangement(n, a));
    assert(count == 24);
    cout << endl;
  }
  {
    int n = 4, k = 2;
    vector<int> a{0, 0};
    cout << endl << n << "^" << k << " arrangements with repeats:" << endl;
    int count = 0;
    do {
      print_range(a.begin(), a.end());
      count++;
      if (count == 13) {
        cout << endl;
      }
    } while (next_arrangement_with_repeats(n, a));
    assert(count == 16);
    cout << endl;
  }
  return 0;
}