cp-book

ecnerwala's competitive programming library

View the Project on GitHub ecnerwala/cp-book

🧪 smawk.test.cpp

View this file on GitHub · Last update: 2026-07-27 01:43:43-07:00

Depends on

Code

#include "smawk.hpp"

#include <catch2/catch_test_macros.hpp>
#include <catch2/catch_get_random_seed.hpp>

#include <algorithm>
#include <random>

struct move_only_t {
	int v;
	move_only_t() : v(-1) {}
	explicit move_only_t(int v_) : v(v_) {
		assert(v_ != -1);
	}
	move_only_t(move_only_t&& o) {
		v = o.v;
		o.v = -1;
	}
	move_only_t& operator = (move_only_t&& o) {
		v = o.v;
		o.v = -1;
		return *this;
	}
	move_only_t(const move_only_t& o) = delete;
	move_only_t& operator = (const move_only_t& o) = delete;
};

std::vector<std::vector<int>> generate_totally_monotone(int N, int M, auto&& rng) {
	std::vector<int> cur_order(M);
	std::iota(cur_order.begin(), cur_order.end(), 0);
	std::vector<int> cur_vals(M);
	std::iota(cur_vals.begin(), cur_vals.end(), 0);
	std::vector<std::vector<int>> perms; perms.reserve(M * (M-1) / 2 + 1);
	perms.push_back(cur_vals);
	{
		std::vector<int> cnds; cnds.reserve(M);
		for (int z = 0; z < M * (M-1) / 2; z++) {
			cnds.clear();
			for (int i = 0; i+1 < M; i++) {
				if (cur_order[i] < cur_order[i+1]) {
					cnds.push_back(i);
				}
			}
			assert(!cnds.empty());
			int i = cnds[std::uniform_int_distribution<int>(0, int(cnds.size()) - 1)(rng)];
			std::swap(cur_order[i], cur_order[i+1]);
			cur_vals[cur_order[i]] = i;
			cur_vals[cur_order[i+1]] = i+1;
			perms.push_back(cur_vals);
		}
	}
	std::vector<std::vector<int>> output; output.reserve(N);
	std::vector<int> stars_bars(M * (M-1) / 2 + N);
	std::fill(stars_bars.begin(), stars_bars.begin() + N, 1);
	std::shuffle(stars_bars.begin(), stars_bars.end(), rng);
	{
		int perm_idx = 0;
		for (auto op : stars_bars) {
			if (op == 0) {
				perm_idx++;
			} else {
				output.push_back(perms[perm_idx]);
			}
		}
	}
	assert(int(output.size()) == N);
	return output;
}

void check_smawk(int N, int M, std::vector<std::vector<int>> mat) {
	auto result = smawk::smawk(N, M, [&](int row, int col) -> move_only_t {
		REQUIRE(0 <= row); REQUIRE(row < N);
		REQUIRE(0 <= col); REQUIRE(col < M);
		return move_only_t{mat[row][col]};
	}, [&](int r, const smawk::value_t<move_only_t>& cnd1, const smawk::value_t<move_only_t>& cnd2) -> bool {
		REQUIRE(0 <= r); REQUIRE(r < N);
		REQUIRE(0 <= cnd1.col); REQUIRE(cnd1.col < M);
		REQUIRE(0 <= cnd2.col); REQUIRE(cnd2.col < M);
		REQUIRE(cnd1.col < cnd2.col);
		REQUIRE(cnd1.v.v == mat[r][cnd1.col]);
		REQUIRE(cnd2.v.v == mat[r][cnd2.col]);
		// cnd2 is better when it's strictly smaller
		return cnd2.v.v < cnd1.v.v;
	});
	REQUIRE(int(result.size()) == N);
	for (int i = 0; i < N; i++) {
		int j = int(std::min_element(mat[i].begin(), mat[i].end()) - mat[i].begin());
		REQUIRE(result[i].col == j);
		REQUIRE(result[i].v.v == mat[i][j]);
	}
}

TEST_CASE("SMAWK", "[smawk]") {
	std::mt19937 mt(Catch::getSeed());
	for (int N : {0, 1, 2, 3, 5, 8, 13}) {
		for (int M : {0, 1, 2, 3, 5, 8, 13}) {
			if (N > 0 && M == 0) continue;
			auto inp = generate_totally_monotone(N, M, mt);
			CAPTURE(N, M, inp);
			check_smawk(N, M, inp);
		}
	}
}

void check_larsch(int N, std::vector<std::vector<int>> mat, bool is_totally_monotone = true) {
	const int M = N;
	smawk::LARSCH l(N, [&](int row, int col) -> move_only_t {
		REQUIRE(0 <= row); REQUIRE(row < N);
		REQUIRE(0 <= col); REQUIRE(col < M);
		REQUIRE(col <= row);
		return move_only_t{mat[row][col]};
	}, [&](int r, const smawk::value_t<move_only_t>& cnd1, const smawk::value_t<move_only_t>& cnd2) -> bool {
		REQUIRE(0 <= r); REQUIRE(r < N);
		REQUIRE(0 <= cnd1.col); REQUIRE(cnd1.col < M);
		REQUIRE(0 <= cnd2.col); REQUIRE(cnd2.col < M);
		REQUIRE(cnd1.col < cnd2.col);
		REQUIRE(cnd1.col <= r);
		REQUIRE(cnd2.col <= r);
		REQUIRE(cnd1.v.v == mat[r][cnd1.col]);
		REQUIRE(cnd2.v.v == mat[r][cnd2.col]);
		// cnd2 is better when it's strictly smaller
		return cnd2.v.v < cnd1.v.v;
	});
	int prv_col = 0;
	for (int i = 0; i < N; i++) {
		auto res = l.push_and_query_next();
		REQUIRE(0 <= res.col); REQUIRE(res.col <= i);
		REQUIRE(res.v.v == mat[i][res.col]);
		if (is_totally_monotone) {
			int j = int(std::min_element(mat[i].begin(), mat[i].begin() + i + 1) - mat[i].begin());
			REQUIRE(res.col == j);
		}
		REQUIRE(res.col >= prv_col);
		prv_col = res.col;
	}
}

TEST_CASE("LARSCH", "[smawk]") {
	std::mt19937 mt(Catch::getSeed());
	for (int N : {0, 1, 2, 3, 5, 8, 13}) {
		auto inp = generate_totally_monotone(N, N, mt);
		for (int i = 0; i < N; i++) {
			for (int j = i+1; j < N; j++) {
				inp[i][j] = -1;
			}
		}
		CAPTURE(N, inp);
		check_larsch(N, inp);
	}
}

TEST_CASE("LARSCH Consistency", "[smawk]") {
	std::mt19937 mt(Catch::getSeed());
	for (int N : {0, 1, 2, 3, 5, 8, 13}) {
		std::vector<std::vector<int>> inp(N);
		for (int i = 0; i < N; i++) {
			inp[i].resize(N);
			std::iota(inp[i].begin(), inp[i].end(), 0);
			std::shuffle(inp[i].begin(), inp[i].end(), mt);
		}
		for (int i = 0; i < N; i++) {
			for (int j = i+1; j < N; j++) {
				inp[i][j] = -1;
			}
		}
		CAPTURE(N, inp);
		// Don't check that it's the global min
		check_larsch(N, inp, false);
	}
}

#include <vector>
#include <cassert>
#include <optional>
#include <concepts>
#include <catch2/catch_test_macros.hpp>
#include <catch2/catch_get_random_seed.hpp>
#include <algorithm>
#include <random>
#line 2 "src/smawk.hpp"

#line 7 "src/smawk.hpp"

namespace smawk {

template <typename T> struct value_t {
	T v;
	int col;
};

// Get(int row, int col) -> T
// Select(int row, const value_t<T>& opt_0, const value_t<T>& opt_1) returns 0 or 1 for which is better
template <typename T, typename Get, typename Select> concept totally_monotone_matrix_oracle =
	std::default_initializable<T> && std::movable<T>
	&& std::invocable<Get, int, int> && std::convertible_to<std::invoke_result_t<Get, int, int>, T>
	&& std::predicate<Select, int, const value_t<T>&, const value_t<T>&>;


template <typename Get, typename Select, typename T = std::invoke_result_t<Get, int, int>>
requires totally_monotone_matrix_oracle<T, Get, Select>
class LARSCH {
public:
	int N;
	Get get;
	Select select;
	int L;
	int num_rows;

	std::vector<std::vector<value_t<T>>> stk;
	std::vector<std::pair<value_t<T>, int>> bests;
	LARSCH() {}
	LARSCH(int N_, Get&& get_, Select&& select_) : N(N_), get(std::forward<Get>(get_)), select(std::forward<Select>(select_)) {
		L = N ? 31 - __builtin_clz(N) : 0;
		stk.resize(L);
		bests.resize(L);
		// N >> L == 1, unless N == 0
		for (int i = 0; i < L; i++) {
			stk[i].reserve(N >> (i+1));
		}
		num_rows = 0;
	}

	value_t<T> push_and_query_next() {
		assert(num_rows < N);
		int inp_row = num_rows++;

		int l = 0;
		value_t<T> nbest;
		while (true) {
			int r = inp_row >> l;
			int col = l == 0 ? inp_row : stk[l-1][r].col;
			if (l == L) {
				// special case: just return the unique element
				assert(r == 0);
				assert(inp_row == 0);
				int row = ((r+1) << l) - 1;
				if (l == 0) nbest = {get(row, col), col};
				else nbest = {std::move(stk[l-1][r].v), stk[l-1][r].col};
				l--;
				break;
			}
			assert(l < L);

			if (r & 1) {
				int row = ((r+1) << l) - 1;
				value_t<T> prv_col_top;
				if (l == 0) prv_col_top = {get(row, col), col};
				else prv_col_top = {std::move(stk[l-1][r].v), stk[l-1][r].col};

				assert(bests[l].second <= r);

				// just check this guy at this row, and then push it into the next layer, but don't query yet
				if (bests[l].second == r || select(row, bests[l].first, prv_col_top)) {
					// prv_col_top is better here
					bests[l].first = std::move(prv_col_top);
					bests[l].second = r;
				}
			}
			if (bests[l].second == r) {
				// We've committed to the new column being the best, so let's prune the stack and propagate to ensure consistency.
				assert(int(stk[l].size()) >= (r+1)/2);
				stk[l].resize((r+1)/2);
				// We can just set .second only, since the only time it's read is the r&1 case, which will override bests[l].first
				if (l+1 < L) bests[l+1].second = (r+1)/2;
			}
			std::optional<value_t<T>> to_push;
			while (int(stk[l].size()) > (r+1)/2) {
				int row = (int(stk[l].size()) << (l+1)) - 1;
				value_t<T> nv{get(row, col), col};
				if (select(row, stk[l].back(), nv)) {
					stk[l].pop_back();
					to_push = std::move(nv);
				} else {
					break;
				}
			}
			if (to_push) {
				stk[l].emplace_back(std::move(*to_push));
			} else {
				int row = (int(stk[l].size()+1) << (l+1)) - 1;
				if (row < N) stk[l].emplace_back(get(row, col), col);
			}
			if (r & 1) {
				// go return
				nbest = std::move(bests[l].first);
				l--;
				break;
			} else if (((r+2) << l) - 1 >= N) {
				// go return
				nbest.col = col;
				break;
			} else {
				l++;
				continue;
			}
			assert(false);
		}
		for (; l >= 0; l--) {
			int r = inp_row >> l;
			assert(!(r & 1));
			int row = ((r+1) << l) - 1;
			bests[l].first = std::move(nbest);
			bool did_set = false;
			while (true) {
				int idx = bests[l].second;
				assert(idx <= r);
				int col = (l == 0 ? idx : stk[l-1][idx].col);
				assert(col <= bests[l].first.col);
				value_t<T> cnd;
				if (l > 0 && idx == r) cnd = {std::move(stk[l-1][r].v), col};
				else cnd = {get(row, col), col};
				if (!did_set || select(row, nbest, cnd)) {
					did_set = true;
					nbest = std::move(cnd);
				}
				if (col == bests[l].first.col) break;
				bests[l].second++;
			}
		}
		assert(l == -1);
		return nbest;
	}
};

template <typename Get, typename Select, typename T = std::invoke_result_t<Get&&, int, int>>
requires totally_monotone_matrix_oracle<T, Get&&, Select&&>
std::vector<value_t<T>> smawk(int N, int M, Get&& get, Select&& select) {
	// TODO: If M >> N, then we should do an extra layer of column filter on the outside. The cutoff should be M > 2N or so.
	std::vector<value_t<T>> res(N);
	for (int i = 0; i < N; i++) res[i].col = -1;
	std::vector<int> stks(N);
	int L = N ? 31 - __builtin_clz(N) : 0;
	std::vector<int> stk_ends(L+1);
	stk_ends[0] = 0;
	for (int l = 0; l < L; l++) {
		int sz = 0;
		auto check_col = [&](int col, int min_sz) -> void {
			while (sz > min_sz) {
				int row = (sz << (l+1)) - 1;
				value_t<T> cnd(get(row, col), col);
				if (select(row, res[row], cnd)) {
					// we prefer cnd, save this
					res[row] = std::move(cnd);
					sz--;
				} else {
					break;
				}
			}

			if (sz < (N >> (l+1))) {
				int row = ((sz+1) << (l+1)) - 1;
				if (res[row].col == col) {
					stks[stk_ends[l] + sz] = col;
					sz++;
				} else {
					value_t<T> cnd(get(row, col), col);
					// This is a legal optimization, but I'm not sure it buys anything real, so just stub it out with true ||
					if (true || res[row].col == -1 || res[row].col < col || !select(row, cnd, res[row])) {
						res[row] = std::move(cnd);
						stks[stk_ends[l] + sz] = col;
						sz++;
					}
				}
			}
		};
		if (l == 0) {
			for (int col = 0; col < M; col++) {
				check_col(col, 0);
			}
		} else {
			for (int z = stk_ends[l-1]; z < stk_ends[l]; z++) {
				check_col(stks[z], (z - stk_ends[l-1]) / 2);
			}
		}
		assert(sz <= (N >> (l+1)));
		stk_ends[l+1] = stk_ends[l] + sz;
	}
	for (int l = L; l >= 0; l--) {
		int z = l == 0 ? 0 : stk_ends[l-1];
		for (int r = 0; r < (N >> l); r += 2) {
			int row = ((r+1) << l) - 1;
			// TODO: You could not reset this? Not sure if it buys anything real.
			res[row].col = -1;
			for (; z < (l == 0 ? M : stk_ends[l]); z++) {
				int col = l == 0 ? z : stks[z];
				value_t<T> cnd = {get(row, col), col};
				if (res[row].col == -1 || select(row, res[row], cnd)) {
					res[row] = std::move(cnd);
				}
				if ((r+1) < (N >> l) && col == res[((r+2) << l) - 1].col) break;
			}
			assert(res[row].col != -1);
		}
	}
	return res;
}

// namespace smawk
}
#line 2 "src/smawk.test.cpp"

#line 5 "src/smawk.test.cpp"

#line 8 "src/smawk.test.cpp"

struct move_only_t {
	int v;
	move_only_t() : v(-1) {}
	explicit move_only_t(int v_) : v(v_) {
		assert(v_ != -1);
	}
	move_only_t(move_only_t&& o) {
		v = o.v;
		o.v = -1;
	}
	move_only_t& operator = (move_only_t&& o) {
		v = o.v;
		o.v = -1;
		return *this;
	}
	move_only_t(const move_only_t& o) = delete;
	move_only_t& operator = (const move_only_t& o) = delete;
};

std::vector<std::vector<int>> generate_totally_monotone(int N, int M, auto&& rng) {
	std::vector<int> cur_order(M);
	std::iota(cur_order.begin(), cur_order.end(), 0);
	std::vector<int> cur_vals(M);
	std::iota(cur_vals.begin(), cur_vals.end(), 0);
	std::vector<std::vector<int>> perms; perms.reserve(M * (M-1) / 2 + 1);
	perms.push_back(cur_vals);
	{
		std::vector<int> cnds; cnds.reserve(M);
		for (int z = 0; z < M * (M-1) / 2; z++) {
			cnds.clear();
			for (int i = 0; i+1 < M; i++) {
				if (cur_order[i] < cur_order[i+1]) {
					cnds.push_back(i);
				}
			}
			assert(!cnds.empty());
			int i = cnds[std::uniform_int_distribution<int>(0, int(cnds.size()) - 1)(rng)];
			std::swap(cur_order[i], cur_order[i+1]);
			cur_vals[cur_order[i]] = i;
			cur_vals[cur_order[i+1]] = i+1;
			perms.push_back(cur_vals);
		}
	}
	std::vector<std::vector<int>> output; output.reserve(N);
	std::vector<int> stars_bars(M * (M-1) / 2 + N);
	std::fill(stars_bars.begin(), stars_bars.begin() + N, 1);
	std::shuffle(stars_bars.begin(), stars_bars.end(), rng);
	{
		int perm_idx = 0;
		for (auto op : stars_bars) {
			if (op == 0) {
				perm_idx++;
			} else {
				output.push_back(perms[perm_idx]);
			}
		}
	}
	assert(int(output.size()) == N);
	return output;
}

void check_smawk(int N, int M, std::vector<std::vector<int>> mat) {
	auto result = smawk::smawk(N, M, [&](int row, int col) -> move_only_t {
		REQUIRE(0 <= row); REQUIRE(row < N);
		REQUIRE(0 <= col); REQUIRE(col < M);
		return move_only_t{mat[row][col]};
	}, [&](int r, const smawk::value_t<move_only_t>& cnd1, const smawk::value_t<move_only_t>& cnd2) -> bool {
		REQUIRE(0 <= r); REQUIRE(r < N);
		REQUIRE(0 <= cnd1.col); REQUIRE(cnd1.col < M);
		REQUIRE(0 <= cnd2.col); REQUIRE(cnd2.col < M);
		REQUIRE(cnd1.col < cnd2.col);
		REQUIRE(cnd1.v.v == mat[r][cnd1.col]);
		REQUIRE(cnd2.v.v == mat[r][cnd2.col]);
		// cnd2 is better when it's strictly smaller
		return cnd2.v.v < cnd1.v.v;
	});
	REQUIRE(int(result.size()) == N);
	for (int i = 0; i < N; i++) {
		int j = int(std::min_element(mat[i].begin(), mat[i].end()) - mat[i].begin());
		REQUIRE(result[i].col == j);
		REQUIRE(result[i].v.v == mat[i][j]);
	}
}

TEST_CASE("SMAWK", "[smawk]") {
	std::mt19937 mt(Catch::getSeed());
	for (int N : {0, 1, 2, 3, 5, 8, 13}) {
		for (int M : {0, 1, 2, 3, 5, 8, 13}) {
			if (N > 0 && M == 0) continue;
			auto inp = generate_totally_monotone(N, M, mt);
			CAPTURE(N, M, inp);
			check_smawk(N, M, inp);
		}
	}
}

void check_larsch(int N, std::vector<std::vector<int>> mat, bool is_totally_monotone = true) {
	const int M = N;
	smawk::LARSCH l(N, [&](int row, int col) -> move_only_t {
		REQUIRE(0 <= row); REQUIRE(row < N);
		REQUIRE(0 <= col); REQUIRE(col < M);
		REQUIRE(col <= row);
		return move_only_t{mat[row][col]};
	}, [&](int r, const smawk::value_t<move_only_t>& cnd1, const smawk::value_t<move_only_t>& cnd2) -> bool {
		REQUIRE(0 <= r); REQUIRE(r < N);
		REQUIRE(0 <= cnd1.col); REQUIRE(cnd1.col < M);
		REQUIRE(0 <= cnd2.col); REQUIRE(cnd2.col < M);
		REQUIRE(cnd1.col < cnd2.col);
		REQUIRE(cnd1.col <= r);
		REQUIRE(cnd2.col <= r);
		REQUIRE(cnd1.v.v == mat[r][cnd1.col]);
		REQUIRE(cnd2.v.v == mat[r][cnd2.col]);
		// cnd2 is better when it's strictly smaller
		return cnd2.v.v < cnd1.v.v;
	});
	int prv_col = 0;
	for (int i = 0; i < N; i++) {
		auto res = l.push_and_query_next();
		REQUIRE(0 <= res.col); REQUIRE(res.col <= i);
		REQUIRE(res.v.v == mat[i][res.col]);
		if (is_totally_monotone) {
			int j = int(std::min_element(mat[i].begin(), mat[i].begin() + i + 1) - mat[i].begin());
			REQUIRE(res.col == j);
		}
		REQUIRE(res.col >= prv_col);
		prv_col = res.col;
	}
}

TEST_CASE("LARSCH", "[smawk]") {
	std::mt19937 mt(Catch::getSeed());
	for (int N : {0, 1, 2, 3, 5, 8, 13}) {
		auto inp = generate_totally_monotone(N, N, mt);
		for (int i = 0; i < N; i++) {
			for (int j = i+1; j < N; j++) {
				inp[i][j] = -1;
			}
		}
		CAPTURE(N, inp);
		check_larsch(N, inp);
	}
}

TEST_CASE("LARSCH Consistency", "[smawk]") {
	std::mt19937 mt(Catch::getSeed());
	for (int N : {0, 1, 2, 3, 5, 8, 13}) {
		std::vector<std::vector<int>> inp(N);
		for (int i = 0; i < N; i++) {
			inp[i].resize(N);
			std::iota(inp[i].begin(), inp[i].end(), 0);
			std::shuffle(inp[i].begin(), inp[i].end(), mt);
		}
		for (int i = 0; i < N; i++) {
			for (int j = i+1; j < N; j++) {
				inp[i][j] = -1;
			}
		}
		CAPTURE(N, inp);
		// Don't check that it's the global min
		check_larsch(N, inp, false);
	}
}

// clang-format off
// @formatter:off
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wpragmas"
#pragma GCC diagnostic ignored "-Wunknown-warning-option"
#pragma GCC diagnostic ignored "-Wmisleading-indentation"
#pragma GCC diagnostic ignored "-Wmultistatement-macros"
#include <bits/stdc++.h>
#include <cassert>
#include <catch2/catch_test_macros.hpp>
#include <catch2/catch_get_random_seed.hpp>
// src/smawk.hpp
namespace smawk{
template<typename T>struct value_t{
T v;
int col;
};
template<typename T,typename Get,typename Select>concept totally_monotone_matrix_oracle=
std::default_initializable<T>&&std::movable<T>
&&std::invocable<Get,int,int>&&std::convertible_to<std::invoke_result_t<Get,int,int>,T>
&&std::predicate<Select,int,const value_t<T>&,const value_t<T>&>;
template<typename Get,typename Select,typename T=std::invoke_result_t<Get,int,int>>
requires totally_monotone_matrix_oracle<T,Get,Select>
class LARSCH{
public:
int N;
Get get;
Select select;
int L;
int num_rows;
std::vector<std::vector<value_t<T>>>stk;
std::vector<std::pair<value_t<T>,int>>bests;
LARSCH(){}
LARSCH(int N_,Get&&get_,Select&&select_):N(N_),get(std::forward<Get>(get_)),select(std::forward<Select>(select_)){
L=N?31-__builtin_clz(N):0;
stk.resize(L);
bests.resize(L);
for(int i=0;i<L;i++){
stk[i].reserve(N>>(i+1));
}
num_rows=0;
}
value_t<T>push_and_query_next(){
assert(num_rows<N);
int inp_row=num_rows++;
int l=0;
value_t<T>nbest;
while(true){
int r=inp_row>>l;
int col=l==0?inp_row:stk[l-1][r].col;
if(l==L){
assert(r==0);
assert(inp_row==0);
int row=((r+1)<<l)-1;
if(l==0)nbest={get(row,col),col};
else nbest={std::move(stk[l-1][r].v),stk[l-1][r].col};
l--;
break;
}
assert(l<L);
if(r&1){
int row=((r+1)<<l)-1;
value_t<T>prv_col_top;
if(l==0)prv_col_top={get(row,col),col};
else prv_col_top={std::move(stk[l-1][r].v),stk[l-1][r].col};
assert(bests[l].second<=r);
if(bests[l].second==r||select(row,bests[l].first,prv_col_top)){
bests[l].first=std::move(prv_col_top);
bests[l].second=r;
}
}
if(bests[l].second==r){
assert(int(stk[l].size())>=(r+1)/2);
stk[l].resize((r+1)/2);
if(l+1<L)bests[l+1].second=(r+1)/2;
}
std::optional<value_t<T>>to_push;
while(int(stk[l].size())>(r+1)/2){
int row=(int(stk[l].size())<<(l+1))-1;
value_t<T>nv{get(row,col),col};
if(select(row,stk[l].back(),nv)){
stk[l].pop_back();
to_push=std::move(nv);
}else{
break;
}
}
if(to_push){
stk[l].emplace_back(std::move(*to_push));
}else{
int row=(int(stk[l].size()+1)<<(l+1))-1;
if(row<N)stk[l].emplace_back(get(row,col),col);
}
if(r&1){
nbest=std::move(bests[l].first);
l--;
break;
}else if(((r+2)<<l)-1>=N){
nbest.col=col;
break;
}else{
l++;
continue;
}
assert(false);
}
for(;l>=0;l--){
int r=inp_row>>l;
assert(!(r&1));
int row=((r+1)<<l)-1;
bests[l].first=std::move(nbest);
bool did_set=false;
while(true){
int idx=bests[l].second;
assert(idx<=r);
int col=(l==0?idx:stk[l-1][idx].col);
assert(col<=bests[l].first.col);
value_t<T>cnd;
if(l>0&&idx==r)cnd={std::move(stk[l-1][r].v),col};
else cnd={get(row,col),col};
if(!did_set||select(row,nbest,cnd)){
did_set=true;
nbest=std::move(cnd);
}
if(col==bests[l].first.col)break;
bests[l].second++;
}
}
assert(l==-1);
return nbest;
}
};
template<typename Get,typename Select,typename T=std::invoke_result_t<Get&&,int,int>>
requires totally_monotone_matrix_oracle<T,Get&&,Select&&>
std::vector<value_t<T>>smawk(int N,int M,Get&&get,Select&&select){
std::vector<value_t<T>>res(N);
for(int i=0;i<N;i++)res[i].col=-1;
std::vector<int>stks(N);
int L=N?31-__builtin_clz(N):0;
std::vector<int>stk_ends(L+1);
stk_ends[0]=0;
for(int l=0;l<L;l++){
int sz=0;
auto check_col=[&](int col,int min_sz)->void{
while(sz>min_sz){
int row=(sz<<(l+1))-1;
value_t<T>cnd(get(row,col),col);
if(select(row,res[row],cnd)){
res[row]=std::move(cnd);
sz--;
}else{
break;
}
}
if(sz<(N>>(l+1))){
int row=((sz+1)<<(l+1))-1;
if(res[row].col==col){
stks[stk_ends[l]+sz]=col;
sz++;
}else{
value_t<T>cnd(get(row,col),col);
if(true||res[row].col==-1||res[row].col<col||!select(row,cnd,res[row])){
res[row]=std::move(cnd);
stks[stk_ends[l]+sz]=col;
sz++;
}
}
}
};
if(l==0){
for(int col=0;col<M;col++){
check_col(col,0);
}
}else{
for(int z=stk_ends[l-1];z<stk_ends[l];z++){
check_col(stks[z],(z-stk_ends[l-1])/2);
}
}
assert(sz<=(N>>(l+1)));
stk_ends[l+1]=stk_ends[l]+sz;
}
for(int l=L;l>=0;l--){
int z=l==0?0:stk_ends[l-1];
for(int r=0;r<(N>>l);r+=2){
int row=((r+1)<<l)-1;
res[row].col=-1;
for(;z<(l==0?M:stk_ends[l]);z++){
int col=l==0?z:stks[z];
value_t<T>cnd={get(row,col),col};
if(res[row].col==-1||select(row,res[row],cnd)){
res[row]=std::move(cnd);
}
if((r+1)<(N>>l)&&col==res[((r+2)<<l)-1].col)break;
}
assert(res[row].col!=-1);
}
}
return res;
}
}
// src/smawk.test.cpp
struct move_only_t{
int v;
move_only_t():v(-1){}
explicit move_only_t(int v_):v(v_){
assert(v_!=-1);
}
move_only_t(move_only_t&&o){
v=o.v;
o.v=-1;
}
move_only_t&operator=(move_only_t&&o){
v=o.v;
o.v=-1;
return*this;
}
move_only_t(const move_only_t&o)=delete;
move_only_t&operator=(const move_only_t&o)=delete;
};
std::vector<std::vector<int>>generate_totally_monotone(int N,int M,auto&&rng){
std::vector<int>cur_order(M);
std::iota(cur_order.begin(),cur_order.end(),0);
std::vector<int>cur_vals(M);
std::iota(cur_vals.begin(),cur_vals.end(),0);
std::vector<std::vector<int>>perms;perms.reserve(M*(M-1)/2+1);
perms.push_back(cur_vals);
{
std::vector<int>cnds;cnds.reserve(M);
for(int z=0;z<M*(M-1)/2;z++){
cnds.clear();
for(int i=0;i+1<M;i++){
if(cur_order[i]<cur_order[i+1]){
cnds.push_back(i);
}
}
assert(!cnds.empty());
int i=cnds[std::uniform_int_distribution<int>(0,int(cnds.size())-1)(rng)];
std::swap(cur_order[i],cur_order[i+1]);
cur_vals[cur_order[i]]=i;
cur_vals[cur_order[i+1]]=i+1;
perms.push_back(cur_vals);
}
}
std::vector<std::vector<int>>output;output.reserve(N);
std::vector<int>stars_bars(M*(M-1)/2+N);
std::fill(stars_bars.begin(),stars_bars.begin()+N,1);
std::shuffle(stars_bars.begin(),stars_bars.end(),rng);
{
int perm_idx=0;
for(auto op:stars_bars){
if(op==0){
perm_idx++;
}else{
output.push_back(perms[perm_idx]);
}
}
}
assert(int(output.size())==N);
return output;
}
void check_smawk(int N,int M,std::vector<std::vector<int>>mat){
auto result=smawk::smawk(N,M,[&](int row,int col)->move_only_t{
REQUIRE(0<=row);REQUIRE(row<N);
REQUIRE(0<=col);REQUIRE(col<M);
return move_only_t{mat[row][col]};
},[&](int r,const smawk::value_t<move_only_t>&cnd1,const smawk::value_t<move_only_t>&cnd2)->bool{
REQUIRE(0<=r);REQUIRE(r<N);
REQUIRE(0<=cnd1.col);REQUIRE(cnd1.col<M);
REQUIRE(0<=cnd2.col);REQUIRE(cnd2.col<M);
REQUIRE(cnd1.col<cnd2.col);
REQUIRE(cnd1.v.v==mat[r][cnd1.col]);
REQUIRE(cnd2.v.v==mat[r][cnd2.col]);
return cnd2.v.v<cnd1.v.v;
});
REQUIRE(int(result.size())==N);
for(int i=0;i<N;i++){
int j=int(std::min_element(mat[i].begin(),mat[i].end())-mat[i].begin());
REQUIRE(result[i].col==j);
REQUIRE(result[i].v.v==mat[i][j]);
}
}
TEST_CASE("SMAWK","[smawk]"){
std::mt19937 mt(Catch::getSeed());
for(int N:{0,1,2,3,5,8,13}){
for(int M:{0,1,2,3,5,8,13}){
if(N>0&&M==0)continue;
auto inp=generate_totally_monotone(N,M,mt);
CAPTURE(N,M,inp);
check_smawk(N,M,inp);
}
}
}
void check_larsch(int N,std::vector<std::vector<int>>mat,bool is_totally_monotone=true){
const int M=N;
smawk::LARSCH l(N,[&](int row,int col)->move_only_t{
REQUIRE(0<=row);REQUIRE(row<N);
REQUIRE(0<=col);REQUIRE(col<M);
REQUIRE(col<=row);
return move_only_t{mat[row][col]};
},[&](int r,const smawk::value_t<move_only_t>&cnd1,const smawk::value_t<move_only_t>&cnd2)->bool{
REQUIRE(0<=r);REQUIRE(r<N);
REQUIRE(0<=cnd1.col);REQUIRE(cnd1.col<M);
REQUIRE(0<=cnd2.col);REQUIRE(cnd2.col<M);
REQUIRE(cnd1.col<cnd2.col);
REQUIRE(cnd1.col<=r);
REQUIRE(cnd2.col<=r);
REQUIRE(cnd1.v.v==mat[r][cnd1.col]);
REQUIRE(cnd2.v.v==mat[r][cnd2.col]);
return cnd2.v.v<cnd1.v.v;
});
int prv_col=0;
for(int i=0;i<N;i++){
auto res=l.push_and_query_next();
REQUIRE(0<=res.col);REQUIRE(res.col<=i);
REQUIRE(res.v.v==mat[i][res.col]);
if(is_totally_monotone){
int j=int(std::min_element(mat[i].begin(),mat[i].begin()+i+1)-mat[i].begin());
REQUIRE(res.col==j);
}
REQUIRE(res.col>=prv_col);
prv_col=res.col;
}
}
TEST_CASE("LARSCH","[smawk]"){
std::mt19937 mt(Catch::getSeed());
for(int N:{0,1,2,3,5,8,13}){
auto inp=generate_totally_monotone(N,N,mt);
for(int i=0;i<N;i++){
for(int j=i+1;j<N;j++){
inp[i][j]=-1;
}
}
CAPTURE(N,inp);
check_larsch(N,inp);
}
}
TEST_CASE("LARSCH Consistency","[smawk]"){
std::mt19937 mt(Catch::getSeed());
for(int N:{0,1,2,3,5,8,13}){
std::vector<std::vector<int>>inp(N);
for(int i=0;i<N;i++){
inp[i].resize(N);
std::iota(inp[i].begin(),inp[i].end(),0);
std::shuffle(inp[i].begin(),inp[i].end(),mt);
}
for(int i=0;i<N;i++){
for(int j=i+1;j<N;j++){
inp[i][j]=-1;
}
}
CAPTURE(N,inp);
check_larsch(N,inp,false);
}
}
#pragma GCC diagnostic pop
// clang-format on
// @formatter:on
Back to top page