enumerate/partition.hpp

View this file on GitHub
Last update: 2025-12-16 20:51:20+09:00
Include: #include "enumerate/partition.hpp"
Link: https://atcoder.jp/contests/abc390/tasks/abc390_d

Depends on

enumerate/bits.hpp

Verified with

Code

#include "enumerate/bits.hpp"

/*
partition は、「減少列」として辞書式の降順に列挙する。
N=10,20,30,40：42, 527, 5604, 37338
N = 50（204226）：12 ms
N = 60（966467）：60 ms
N = 70（4087968）：270 ms
N = 80（15796476）：1100 ms
N = 90（56634173）：4800 ms
N = 100 (190569292) : 15600 ms
*/
template <typename F>
void enumerate_partition(int N, F query, int LIM_len = -1, int LIM_val = -1) {
  assert(N >= 0);
  auto dfs = [&](auto self, vc<int> &p, int sum) -> void {
    if (sum == N) {
      query(p);
      return;
    }
    if (LIM_len != -1 && len(p) == LIM_len) return;
    int nxt = (len(p) == 0 ? N : p.back());
    if (LIM_val != -1) chmin(nxt, LIM_val);
    chmin(nxt, N - sum);
    p.eb(0);
    FOR3_R(x, 1, nxt + 1) {
      p.back() = x;
      self(self, p, sum + x);
    }
    p.pop_back();
  };
  vc<int> p;
  dfs(dfs, p, 0);
}

// N 元集合の分割の列挙 (Bell number)
// f({s0,s1,...}), f(vc<int>)
// https://atcoder.jp/contests/abc390/tasks/abc390_d
// N = 11（678570）：29 ms
// N = 12（4213597）：208 ms
// N = 13（27644437）：2084 ms
template <typename F>
void enumerate_set_partition(int N, F f) {
  vc<u32> S;
  auto dfs = [&](auto &dfs, u32 rest) -> void {
    if (rest == 0) {
      return f(S);
    }
    int a = lowbit(rest);
    rest -= u32(1) << a;
    enumerate_all_subset<u32, true>(rest, [&](u32 s) -> void {
      S.eb(s | 1 << a);
      dfs(dfs, rest - s);
      POP(S);
    });
  };
  dfs(dfs, (u32(1) << N) - 1);
}

#line 1 "enumerate/bits.hpp"
template <typename BS, typename F>
void enumerate_bits_bitset(BS& b, int L, int R, F&& f) {
  if (L >= len(b)) return;
  int p = (b[L] ? L : b._Find_next(L));
  while (p < R) {
    f(p);
    p = b._Find_next(p);
  }
}

template <typename UINT, typename F>
inline void enumerate_all_bit(UINT s, F&& f) {
  static_assert(is_unsigned<UINT>::value);
  while (s) {
    f(lowbit(s));
    s &= s - 1;
  }
}

template <typename UINT, bool inc_empty, typename F>
inline void enumerate_all_subset(UINT s, F&& f) {
  static_assert(is_unsigned<UINT>::value);
  for (UINT t = s; t; t = (t - 1) & s) f(t);
  if constexpr (inc_empty) f(0);
}
#line 2 "enumerate/partition.hpp"

/*
partition は、「減少列」として辞書式の降順に列挙する。
N=10,20,30,40：42, 527, 5604, 37338
N = 50（204226）：12 ms
N = 60（966467）：60 ms
N = 70（4087968）：270 ms
N = 80（15796476）：1100 ms
N = 90（56634173）：4800 ms
N = 100 (190569292) : 15600 ms
*/
template <typename F>
void enumerate_partition(int N, F query, int LIM_len = -1, int LIM_val = -1) {
  assert(N >= 0);
  auto dfs = [&](auto self, vc<int> &p, int sum) -> void {
    if (sum == N) {
      query(p);
      return;
    }
    if (LIM_len != -1 && len(p) == LIM_len) return;
    int nxt = (len(p) == 0 ? N : p.back());
    if (LIM_val != -1) chmin(nxt, LIM_val);
    chmin(nxt, N - sum);
    p.eb(0);
    FOR3_R(x, 1, nxt + 1) {
      p.back() = x;
      self(self, p, sum + x);
    }
    p.pop_back();
  };
  vc<int> p;
  dfs(dfs, p, 0);
}

// N 元集合の分割の列挙 (Bell number)
// f({s0,s1,...}), f(vc<int>)
// https://atcoder.jp/contests/abc390/tasks/abc390_d
// N = 11（678570）：29 ms
// N = 12（4213597）：208 ms
// N = 13（27644437）：2084 ms
template <typename F>
void enumerate_set_partition(int N, F f) {
  vc<u32> S;
  auto dfs = [&](auto &dfs, u32 rest) -> void {
    if (rest == 0) {
      return f(S);
    }
    int a = lowbit(rest);
    rest -= u32(1) << a;
    enumerate_all_subset<u32, true>(rest, [&](u32 s) -> void {
      S.eb(s | 1 << a);
      dfs(dfs, rest - s);
      POP(S);
    });
  };
  dfs(dfs, (u32(1) << N) - 1);
}