nt/EGZ.hpp

• View this file on GitHub
• Last update: 2026-08-31 22:36:55+09:00
• Include: #include "nt/EGZ.hpp"

Depends on

• ds/csr.hpp
• mod/mod_inv.hpp

Verified with

• test/1_mytest/EGZ.test.cpp

Code

#include "ds/csr.hpp"
#include "mod/mod_inv.hpp"

// p-subset で総和が 0 mod p のものを作る
// return: indices
vc<int> EGZ_prime(int p, vc<ll> A) {
  assert(len(A) == p + p - 1);
  for (auto& x : A) x = bmod<ll>(x, p);
  CSR<int> ids(p);
  FOR(i, len(A)) { ids.add(A[i], i); }
  ids.build();

  A.clear();
  FOR(x, p) FOR(len(ids[x])) A.eb(x);

  [&]() -> void {
    FOR(i, p) {
      if (A[i] == A[i + p - 1]) {
        A = {A.begin() + i, A.begin() + i + p};
        return;
      }
    }
    int t = 0;
    FOR(i, p) t = (t + p - A[i]) % p;
    vc<int> par(p, -1);
    auto exist = [&](int i) -> bool { return (i == 0 || par[i] != -1); };
    FOR(i, p - 1) {
      if (exist(t)) break;
      int d = A[i + p] - A[i];
      ll L = 0, R = mod_inv(d, p) * t % p;
      while (L + 1 < R) {
        ll M = (L + R) / 2;
        (exist(M * d % p) ? L : R) = M;
      }
      par[R * d % p] = i;
    }
    while (t != 0) {
      int i = par[t];
      int d = A[i + p] - A[i];
      t = (t + p - d) % p;
      A[i] = A[i + p];
    }
    A.resize(p);
  }();
  vc<int> CNT(p);
  for (auto& x : A) CNT[x]++;
  vc<int> res;
  FOR(x, p) {
    for (int i : ids[x]) {
      if (CNT[x]) --CNT[x], res.eb(i);
    }
  }
  return res;
}

// N-subset で総和が 0 mod p のものを作る
// return: indices
vc<int> EGZ(int N, vc<ll> A) {
  for (auto& x : A) x = bmod<ll>(x, N);
  assert(len(A) == 2 * N - 1);
  if (N == 1) return {0};
  int p = 2;
  while (N % p != 0) ++p;
  if (N == p) return EGZ_prime(N, A);

  // p is a prime factor
  int M = N / p;
  vc<int> ids;
  vc<int> yet;
  vi nxt_val;
  int k = 0;
  // p-EGZ * (2M-1)
  vc<int> used(2 * p - 1);
  FOR(2 * M - 1) {
    while (len(yet) < 2 * p - 1) {
      yet.eb(k++);
    }
    vc<ll> B = rearrange(A, yet);
    vc<int> way = EGZ_prime(p, B);
    FOR(i, 2 * p - 1) used[i] = 0;
    for (int i : way) used[i] = 1;
    vc<int> nxt;
    ll x = 0;
    FOR(i, 2 * p - 1) {
      if (used[i]) {
        x += A[yet[i]];
        ids.eb(yet[i]);
      } else {
        nxt.eb(yet[i]);
      }
    }
    swap(yet, nxt);
    assert(x % p == 0);
    nxt_val.eb(x / p);
  }
  vc<int> I = EGZ(M, nxt_val);
  vc<int> res;
  for (int i : I) {
    FOR(j, p * i, p * i + p) res.eb(ids[j]);
  }
  return res;
}

#line 1 "ds/csr.hpp"

template <typename T>
struct CSR {
  int n;
  bool prepared;
  vc<int> ptr;
  vc<int> I;
  vc<T> dat;

  CSR(int n = 0) : n(n), prepared(false) {}
  void reserve(int n) { dat.reserve(n); }

  void add(int i, const T& x) {
    assert(0 <= i && i < n && !prepared);
    I.eb(i), dat.eb(x);
  }

  void build() {
    assert(!prepared);
    prepared = 1;
    ptr.assign(n + 1, 0);
    for (auto& i : I) ptr[1 + i]++;
    FOR(i, len(ptr) - 1) ptr[i + 1] += ptr[i];
    vc<T> tmp(len(dat));
    FOR(k, len(dat)) {
      int i = I[k];
      tmp[ptr[i]++] = dat[k];
    }
    swap(dat, tmp);
    ptr.pop_back();
    ptr.insert(ptr.begin(), 0);
    I.clear();
  }

  struct range {
    T *first, *last;
    T* begin() const { return first; }
    T* end() const { return last; }
    bool empty() const { return first == last; }
    int size() const { return last - first; }
  };

  range operator[](int i) {
    assert(prepared);
    return range{dat.data() + ptr[i], dat.data() + ptr[i + 1]};
  }
};
#line 1 "mod/mod_inv.hpp"

// long でも大丈夫
// (val * x - 1) が mod の倍数になるようにする
// 特に mod=0 なら x=0 が満たす
ll mod_inv(ll val, ll mod) {
  if (mod == 0) return 0;
  mod = abs(mod);
  val %= mod;
  if (val < 0) val += mod;
  ll a = val, b = mod, u = 1, v = 0, t;
  while (b > 0) {
    t = a / b;
    swap(a -= t * b, b), swap(u -= t * v, v);
  }
  if (u < 0) u += mod;
  return u;
}
#line 3 "nt/EGZ.hpp"

// p-subset で総和が 0 mod p のものを作る
// return: indices
vc<int> EGZ_prime(int p, vc<ll> A) {
  assert(len(A) == p + p - 1);
  for (auto& x : A) x = bmod<ll>(x, p);
  CSR<int> ids(p);
  FOR(i, len(A)) { ids.add(A[i], i); }
  ids.build();

  A.clear();
  FOR(x, p) FOR(len(ids[x])) A.eb(x);

  [&]() -> void {
    FOR(i, p) {
      if (A[i] == A[i + p - 1]) {
        A = {A.begin() + i, A.begin() + i + p};
        return;
      }
    }
    int t = 0;
    FOR(i, p) t = (t + p - A[i]) % p;
    vc<int> par(p, -1);
    auto exist = [&](int i) -> bool { return (i == 0 || par[i] != -1); };
    FOR(i, p - 1) {
      if (exist(t)) break;
      int d = A[i + p] - A[i];
      ll L = 0, R = mod_inv(d, p) * t % p;
      while (L + 1 < R) {
        ll M = (L + R) / 2;
        (exist(M * d % p) ? L : R) = M;
      }
      par[R * d % p] = i;
    }
    while (t != 0) {
      int i = par[t];
      int d = A[i + p] - A[i];
      t = (t + p - d) % p;
      A[i] = A[i + p];
    }
    A.resize(p);
  }();
  vc<int> CNT(p);
  for (auto& x : A) CNT[x]++;
  vc<int> res;
  FOR(x, p) {
    for (int i : ids[x]) {
      if (CNT[x]) --CNT[x], res.eb(i);
    }
  }
  return res;
}

// N-subset で総和が 0 mod p のものを作る
// return: indices
vc<int> EGZ(int N, vc<ll> A) {
  for (auto& x : A) x = bmod<ll>(x, N);
  assert(len(A) == 2 * N - 1);
  if (N == 1) return {0};
  int p = 2;
  while (N % p != 0) ++p;
  if (N == p) return EGZ_prime(N, A);

  // p is a prime factor
  int M = N / p;
  vc<int> ids;
  vc<int> yet;
  vi nxt_val;
  int k = 0;
  // p-EGZ * (2M-1)
  vc<int> used(2 * p - 1);
  FOR(2 * M - 1) {
    while (len(yet) < 2 * p - 1) {
      yet.eb(k++);
    }
    vc<ll> B = rearrange(A, yet);
    vc<int> way = EGZ_prime(p, B);
    FOR(i, 2 * p - 1) used[i] = 0;
    for (int i : way) used[i] = 1;
    vc<int> nxt;
    ll x = 0;
    FOR(i, 2 * p - 1) {
      if (used[i]) {
        x += A[yet[i]];
        ids.eb(yet[i]);
      } else {
        nxt.eb(yet[i]);
      }
    }
    swap(yet, nxt);
    assert(x % p == 0);
    nxt_val.eb(x / p);
  }
  vc<int> I = EGZ(M, nxt_val);
  vc<int> res;
  for (int i : I) {
    FOR(j, p * i, p * i + p) res.eb(ids[j]);
  }
  return res;
}

Back to top page