6.2.5 Enumerating Partitions

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

/*

A partition of a natural number $n$ is a way to write $n$ as a sum of positive integers where the
order of the addends does not matter.

The lexicographic successor edits the suffix of a non-increasing partition: increase the rightmost
part that can grow without breaking the order, then refill the remaining sum with ones. Ranking and
unranking use a dynamic-programming table that counts partitions of a remaining sum by their next
part. This lets the algorithms skip whole blocks of partitions sharing a common prefix.

- `next_partition(p)` takes a reference to a vector `p` of positive integers as a partition of $n$
  for which the function will re-assign to become the next lexicographically greater partition. The
  function returns true if such a partition exists, or false if `p` already consists of the
  lexicographically greatest partition (i.e. the single integer $n$).
- `partition_by_rank(n, r)` returns the partition of $n$ that is lexicographically ranked $r$ if
  addends in each partition were sorted in non-increasing order, where $r$ is a 0-based rank in the
  range $[0, p(n))$ where $p(n)$ is the number of partitions of $n$.
- `rank_by_partition(p)` returns an integer representing the 0-based rank of the partition specified
  by vector `p`, which must consist of positive integers sorted in non-increasing order.
- `generate_increasing_partitions(n, f)` calls the function `f(lo, hi)` on strictly increasing
  partitions of $n$ in lexicographically increasing order of partition, where `lo` and `hi` are
  random-access iterators to a range $[`lo`, `hi`)$ of integers. Note that non-strictly increasing
  partitions like $\{1, 1, 1, 1\}$ are skipped.

Time Complexity:
- O(n) per call to `next_partition()`.
- O(n^2) per call to `partition_by_rank()` and `rank_by_partition()`.
- O(p(n)) per call to `generate_increasing_partitions()`, where $p(n)$ is the number of partitions
  of $n$.

Space Complexity:
- O(1) auxiliary for `next_partition()`.
- O(n^2) auxiliary heap space for `partition_function()`, `partition_by_rank()`, and
  `rank_by_partition()`.
- O(n) auxiliary stack space for `generate_increasing_partitions()`.

*/

#include <cstdint>
#include <vector>

bool next_partition(std::vector<int> &p) {
  int n = static_cast<int>(p.size());
  if (n <= 1) {
    return false;
  }
  int s = p[n - 1] - 1, i = n - 2;
  p.pop_back();
  for (; i > 0 && p[i] == p[i - 1]; i--) {
    s += p[i];
    p.pop_back();
  }
  for (p[i]++; s > 0; s--) {
    p.push_back(1);
  }
  return true;
}

int64_t partition_function(int a, int b) {
  static std::vector<std::vector<int64_t>> p(1, std::vector<int64_t>(1, 1));
  if (a >= static_cast<int>(p.size())) {
    int old = static_cast<int>(p.size());
    p.resize(a + 1);
    for (int i = 0; i < old; i++) {
      p[i].resize(a + 1);
    }
    for (int i = old; i <= a; i++) {
      p[i].resize(a + 1);
      for (int j = 1; j <= i; j++) {
        p[i][j] = p[i - 1][j - 1] + p[i - j][j];
      }
    }
  }
  return p[a][b];
}

std::vector<int> partition_by_rank(int n, int64_t r) {
  std::vector<int> res;
  for (int i = n, j; i > 0; i -= j) {
    for (j = 1;; j++) {
      int64_t count = partition_function(i, j);
      if (r < count) {
        break;
      }
      r -= count;
    }
    res.push_back(j);
  }
  return res;
}

int64_t rank_by_partition(const std::vector<int> &p) {
  int64_t res = 0;
  int sum = 0;
  for (int x : p) {
    sum += x;
  }
  for (int x : p) {
    for (int j = 0; j < x; j++) {
      res += partition_function(sum, j);
    }
    sum -= x;
  }
  return res;
}

template<typename Fn>
void generate_increasing_partitions(int left, int prev, int i, std::vector<int> &p, Fn f) {
  if (left == 0) {
    f(p.begin(), p.begin() + i);
    return;
  }
  for (p[i] = prev + 1; p[i] <= left; p[i]++) {
    generate_increasing_partitions(left - p[i], p[i], i + 1, p, f);
  }
}

template<typename Fn>
void generate_increasing_partitions(int n, Fn f) {
  std::vector<int> p(n);
  generate_increasing_partitions(n, 0, 0, p, f);
}

/*** Example Usage and Output:

Partitions of 4:
{1,1,1,1} {2,1,1} {2,2} {3,1} {4}

Increasing partitions of 8:
{1,2,5} {1,3,4} {1,7} {2,6} {3,5} {8}

***/

#include <algorithm>
#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() {
  // Exercise incremental growth of the cached dynamic-programming table.
  assert(partition_function(3, 1) == 1);
  assert(partition_function(8, 3) == 5);
  {
    int n = 4;
    vector<int> a(n, 1);
    cout << "Partitions of " << n << ":" << endl;
    int count = 0;
    do {
      print_range(a.begin(), a.end());
      assert(a == partition_by_rank(n, count));
      assert(rank_by_partition(a) == count);
      count++;
    } while (next_partition(a));
    assert(count == 5);
    cout << endl;
  }
  {
    int n = 8;
    cout << "\nIncreasing partitions of " << n << ":" << endl;
    int count = 0;
    generate_increasing_partitions(n, [&](auto lo, auto hi) {
      print_range(lo, hi);
      count++;
    });
    assert(count == 6);
    cout << endl;
  }
  return 0;
}