区間制約付き一次不定方程式

math/linear_diophantine.hpp · GitHub で見る

共通ページ

\[ax+by=c,\qquad x_l\le x<x_r,\qquad y_l\le y<y_r\]

を満たす整数の組 \((x,y)\) を求めます。係数と区間端点は負でも構いません。解を列挙する配列を持たずに、個数と辞書順で指定した解を取得できます。

#include "lib/math/linear_diophantine.hpp"
#include <cassert>
#include <cstdint>
#include <utility>

const auto a = linear_diophantine(2, 3, 12, 0, 7, 0, 5);
assert(a && a->count() == 3);
assert((a->first() == std::pair<std::int64_t, std::int64_t>(0, 4)));
assert((a->kth(1) == std::pair<std::int64_t, std::int64_t>(3, 2)));
assert((a->kth(2) == std::pair<std::int64_t, std::int64_t>(6, 0)));
assert(!linear_diophantine(2, 4, 1, -10, 10, -10, 10));

const auto rectangle = linear_diophantine(0, 0, 0, 2, 5, 7, 9);
assert(rectangle && rectangle->count() == 6);
assert((rectangle->kth(3) == std::pair<std::int64_t, std::int64_t>(3, 8)));

操作・型	引数・返り値
linear_diophantine(a, b, c, xl, xr, yl, yr)	7 引数は全て std::int64_t。std::optional<LinearDiophantineSolution> を返す。区間は [xl, xr) と [yl, yr)
solution.count()	解の個数を __uint128_t で返す
solution.first()	辞書順最小の解を std::pair<std::int64_t, std::int64_t> で返す
solution.kth(k)	__uint128_t 型の 0 始まりの順位 k にある解を std::pair<std::int64_t, std::int64_t> で返す。x、次に y の昇順

全引数は符号付き 64 bit 整数の全範囲に対応します。区間が空、方程式に整数解がない、または解が区間内にない場合は std::nullopt を返します。xl > xr または yl > yr は std::invalid_argument、k >= count() は std::out_of_range を送出します。負の順位は指定しないでください。既定引数はなく、全操作が constexpr に対応します。

a = b = 0 のときも扱えます。c = 0 なら区間内の全ての組が解で、それ以外は解なしです。解の個数は 64 bit を超え得るため __uint128_t で返します。この型は有効な全入力の個数を表現できます。座標の返り値は 64 bit、途中の計算は 128 bit 整数を使います。

係数の絶対値の最大を \(A\) として、構築は \(O(\log A)\) 時間・\(O(1)\) 追加空間です。全ての係数が 0 の場合は \(O(1)\) 時間です。count()・first()・kth(k) は \(O(1)\) 時間です。

依存ライブラリ

• 整数演算の共通部品
• 整数除算の切り下げ・切り上げ

コード

#ifndef LINEAR_DIOPHANTINE_HPP
#define LINEAR_DIOPHANTINE_HPP

#include "integer_division.hpp"
#include <algorithm>
#include <cstdint>
#include <limits>
#include <optional>
#include <stdexcept>
#include <utility>

class LinearDiophantineSolution {
	std::int64_t x_, y_;
	__int128_t dx_, dy_;
	__uint128_t count_, width_;
	constexpr LinearDiophantineSolution(std::int64_t x, std::int64_t y, __int128_t dx, __int128_t dy, __uint128_t count, __uint128_t width = 0)
		: x_(x), y_(y), dx_(dx), dy_(dy), count_(count), width_(width) {}
	friend constexpr std::optional<LinearDiophantineSolution> linear_diophantine(std::int64_t, std::int64_t, std::int64_t, std::int64_t, std::int64_t, std::int64_t, std::int64_t);

  public:
	constexpr __uint128_t count() const { return count_; }
	constexpr std::pair<std::int64_t, std::int64_t> kth(__uint128_t k) const {
		if (k >= count_) throw std::out_of_range("linear_diophantine index");
		if (width_) return {std::int64_t(__int128_t(x_) + __int128_t(k / width_)), std::int64_t(__int128_t(y_) + __int128_t(k % width_))};
		return {std::int64_t(x_ + __int128_t(k) * dx_), std::int64_t(y_ + __int128_t(k) * dy_)};
	}
	constexpr std::pair<std::int64_t, std::int64_t> first() const { return {x_, y_}; }
};

constexpr std::optional<LinearDiophantineSolution> linear_diophantine(std::int64_t a, std::int64_t b, std::int64_t c, std::int64_t xl, std::int64_t xr, std::int64_t yl, std::int64_t yr) {
	using I = __int128_t;
	if (xl > xr || yl > yr) throw std::invalid_argument("linear_diophantine interval");
	if (xl == xr || yl == yr) return std::nullopt;
	if (!a && !b) {
		if (c) return std::nullopt;
		__uint128_t width = I(yr) - yl;
		return LinearDiophantineSolution(xl, yl, 0, 0, __uint128_t(I(xr) - xl) * width, width);
	}
	std::uint64_t r = a < 0 ? -I(a) : I(a), s = b < 0 ? -I(b) : I(b);
	I x = 1, u = 0, y = 0, v = 1;
	if (s) for (;;) {
		std::uint64_t q = r / s;
		r -= q * s;
		if (!r) {
			r = s, x = u, y = v;
			break;
		}
		x -= q * u;
		y -= q * v;
		std::swap(r, s);
		std::swap(x, u);
		std::swap(y, v);
	}
	if (I(c) % r) return std::nullopt;
	x *= I(c) / r;
	y *= I(c) / r;
	if (a < 0) x = -x;
	if (b < 0) y = -y;
	I dx = I(b) / r, dy = -I(a) / r;
	if (dx < 0 || (!dx && dy < 0)) dx = -dx, dy = -dy;
	I left = std::numeric_limits<I>::min(), right = std::numeric_limits<I>::max();
	auto restrict = [&](I origin, I step, I lo, I hi) {
		if (!step) return lo <= origin && origin < hi;
		I l = step > 0 ? ceil_div(lo - origin, step) : ceil_div(hi - 1 - origin, step);
		I r = step > 0 ? floor_div(hi - 1 - origin, step) : floor_div(lo - origin, step);
		left = std::max(left, l);
		right = std::min(right, r);
		return left <= right;
	};
	if (!restrict(x, dx, xl, xr) || !restrict(y, dy, yl, yr)) return std::nullopt;
	return LinearDiophantineSolution(std::int64_t(x + left * dx), std::int64_t(y + left * dy), dx, dy, __uint128_t(right - left) + 1);
}

#endif

Back to top page