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

/*

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.

- `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, \text{partitions}(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(n, r)` and `rank_by_partition(p)`.
- O(p(n)) per call to `generate_increasing_partitions(n, f)`, 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);
    p[0].resize(a + 1);
    for (int i = 1; i <= a; i++) {
      p[i].resize(a + 1);
      for (int j = old; 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;
}

using Fn = void (*)(std::vector<int>::iterator, std::vector<int>::iterator);

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);
  }
}

void generate_increasing_partitions(int n, Fn f) {
  std::vector<int> p(n, 0);
  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 <cassert>
#include <iostream>
using namespace std;

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

int main() {
  {
    int n = 4;
    vector<int> a(n, 1);
    cout << "Partitions of " << n << ":" << endl;
    int count = 0;
    do {
      print_range(a.begin(), a.end());
      vector<int> b = partition_by_rank(n, count);
      assert(equal(a.begin(), a.end(), b.begin()));
      assert(rank_by_partition(a) == count);
      count++;
    } while (next_partition(a));
    cout << endl;
  }
  {
    int n = 8;
    cout << "\nIncreasing partitions of " << n << ":" << endl;
    generate_increasing_partitions(n, print_range);
    cout << endl;
  }
  return 0;
}