m1une's library

This documentation is automatically generated by online-judge-tools/verification-helper

View on GitHub

:heavy_check_mark: Beam Search
(heuristic/beam_search.hpp)

What It Does

Beam search constructs solutions one decision at a time. At each depth:

  1. every state in the current beam generates children;
  2. every child receives a score;
  3. only the best beam_width children survive.

Unlike local search, beam search holds several states at once. Unlike exhaustive search, its memory stays bounded. A larger beam usually explores more choices but costs more time and memory.

The easiest interface lets expand return a vector of children:

auto expand = [](const State& state, int next_depth) {
    std::vector<State> children;
    // Add every possible one-step continuation to children.
    return children;
};

next_depth is 1 while expanding the initial state, then 2, and so on. An overload without the depth parameter is also accepted:

auto expand = [](const State& state) {
    std::vector<State> children;
    return children;
};

Complete Example

This example chooses exactly three distinct values while maximizing their sum. The state stores which index may be chosen next, the chosen indices, and the incrementally maintained score.

#include "heuristic/beam_search.hpp"

#include <iostream>
#include <utility>
#include <vector>

struct State {
    int next_index = 0;
    std::vector<int> chosen_indices;
    int score = 0;
};

int main() {
    const std::vector<int> values = {7, -2, 10, 5, 8, 4};

    auto expand = [&](const State& state, int) {
        std::vector<State> children;
        for (int index = state.next_index;
             index < int(values.size()); index++) {
            State child = state;
            child.next_index = index + 1;
            child.chosen_indices.push_back(index);
            child.score += values[index];
            children.push_back(std::move(child));
        }
        return children;
    };

    auto evaluate = [](const State& state) {
        return state.score;
    };

    State initial;
    auto result = m1une::heuristic::beam_search(
        std::move(initial), 3, 100, expand, evaluate,
        m1une::heuristic::Objective::maximize);

    std::cout << result.score << '\n';
    for (int index : result.state.chosen_indices) {
        std::cout << index << ' ';
    }
    std::cout << '\n';
}

Here depth_limit = 3 means three indices are selected. beam_width = 100 is large enough to retain every partial combination in this small example. On a real problem, reduce the width until time and memory fit the contest limit.

Function Contract

template <class State, class Expand, class Evaluate>
auto beam_search(
    State initial_state,
    int depth_limit,
    int beam_width,
    Expand expand,
    Evaluate evaluate,
    Objective objective = Objective::maximize
);
Argument Meaning
initial_state State at depth zero. It is moved into the search.
depth_limit Maximum number of construction steps; must be nonnegative.
beam_width Maximum states retained after each step; must be positive.
expand Returns or emits all one-step children of a state.
evaluate Returns the score used to rank states at the same depth.
objective Chooses whether smaller or larger scores are better.

The returned type is BeamSearchResult<State, Score>:

Member Meaning
State state Best state in the last non-empty layer.
Score score evaluate(state).
int depth Depth of the returned layer.
size_t expanded_states Number of calls to expand.
size_t generated_states Number of children generated before pruning.

If no state generates a child, search stops early and returns the best state in the previous non-empty layer. The function returns the best state from the final layer, not the best score seen across incomparable depths.

If C_d children are generated at depth d and B is the beam width, time is $O(\sum_d C_d \log(B+1))$ plus expansion and evaluation. Extra memory is $O(B)$ states, excluding temporary vectors returned by the simple interface.

Allocation-Conscious Expansion

Returning a vector is easiest, but it temporarily stores all children of one state. For large branching factors, use the emitter interface. Instead of returning a vector, accept a third argument and call it once per child:

auto expand = [&](const State& state, int next_depth, auto emit) {
    for (int move = 0; move < move_count(state); move++) {
        State child = make_child(state, move, next_depth);
        emit(std::move(child));
    }
};

The call to beam_search is otherwise identical. Each emitted state is scored immediately, and rejected states are discarded without building a child vector. This form also supports move-only state types.

Designing the Score

evaluate ranks partial states at the same depth. A useful beam score often combines the value already obtained with an estimate of future potential. For maximization, a typical pattern is:

beam_score = current_value + estimated_remaining_value

A final-answer score alone may rank partial states poorly. If a promising path starts slowly, a narrow beam can delete it before its later benefit appears. Improving the estimate or increasing beam_width can help.

Common Mistakes

Depends on

Required by

Verified with

Code

#ifndef M1UNE_HEURISTIC_BEAM_SEARCH_HPP
#define M1UNE_HEURISTIC_BEAM_SEARCH_HPP 1

#include <algorithm>
#include <cassert>
#include <concepts>
#include <cstddef>
#include <functional>
#include <type_traits>
#include <utility>
#include <vector>

#include "objective.hpp"

namespace m1une {
namespace heuristic {

template <class State, class Score>
struct BeamSearchResult {
    State state;
    Score score;
    int depth;
    std::size_t expanded_states;
    std::size_t generated_states;
};

namespace beam_search_detail {

template <class State, class Score>
struct Node {
    State state;
    Score score;
    std::size_t order;
};

template <class State, class Score>
struct BetterNode {
    Objective objective;

    bool operator()(const Node<State, Score>& first,
                    const Node<State, Score>& second) const {
        if (better_score(first.score, second.score, objective)) return true;
        if (better_score(second.score, first.score, objective)) return false;
        return first.order < second.order;
    }
};

}  // namespace beam_search_detail

// expand(state, next_depth) may return a range of children. For allocation-free
// generation, expand(state, next_depth, emit) may instead call emit(child).
// evaluate(state) returns its score. The best beam_width states are retained at
// every depth, and the best state in the last non-empty layer is returned.
template <class State, class Expand, class Evaluate>
auto beam_search(State initial_state, int depth_limit, int beam_width,
                 Expand expand, Evaluate evaluate,
                 Objective objective = Objective::maximize) {
    assert(0 <= depth_limit);
    assert(0 < beam_width);

    using Score = std::remove_cvref_t<
        std::invoke_result_t<Evaluate&, const State&>>;
    using Node = beam_search_detail::Node<State, Score>;
    using Better = beam_search_detail::BetterNode<State, Score>;

    Score initial_score = std::invoke(evaluate, initial_state);
    std::vector<Node> beam;
    beam.push_back(Node{std::move(initial_state),
                        std::move(initial_score), 0});

    std::size_t expanded_states = 0;
    std::size_t generated_states = 0;
    int reached_depth = 0;
    if (depth_limit < 0 || beam_width <= 0) depth_limit = 0;

    Better better{objective};
    for (int next_depth = 1; next_depth <= depth_limit; next_depth++) {
        std::vector<Node> candidates;
        candidates.reserve(static_cast<std::size_t>(beam_width));
        std::size_t order = 0;

        for (const Node& node : beam) {
            expanded_states++;
            auto emit = [&](auto&& candidate_state) {
                using Candidate = decltype(candidate_state);
                static_assert(std::is_constructible_v<State, Candidate>);
                State state(std::forward<Candidate>(candidate_state));
                Score candidate_score = std::invoke(evaluate, state);
                Node candidate{std::move(state), std::move(candidate_score),
                               order++};
                generated_states++;
                if (int(candidates.size()) < beam_width) {
                    candidates.push_back(std::move(candidate));
                    std::push_heap(candidates.begin(), candidates.end(), better);
                } else if (better(candidate, candidates.front())) {
                    std::pop_heap(candidates.begin(), candidates.end(), better);
                    candidates.back() = std::move(candidate);
                    std::push_heap(candidates.begin(), candidates.end(), better);
                }
            };
            if constexpr (std::invocable<Expand&, const State&, int>) {
                auto next_states =
                    std::invoke(expand, node.state, next_depth);
                for (auto& candidate_state : next_states) {
                    emit(std::move(candidate_state));
                }
            } else if constexpr (std::invocable<Expand&, const State&>) {
                auto next_states = std::invoke(expand, node.state);
                for (auto& candidate_state : next_states) {
                    emit(std::move(candidate_state));
                }
            } else {
                std::invoke(expand, node.state, next_depth, emit);
            }
        }

        if (candidates.empty()) break;
        beam = std::move(candidates);
        reached_depth = next_depth;
    }

    int best = 0;
    for (int index = 1; index < int(beam.size()); index++) {
        if (better(beam[index], beam[best])) best = index;
    }
    return BeamSearchResult<State, Score>{
        std::move(beam[best].state), std::move(beam[best].score),
        reached_depth, expanded_states, generated_states};
}

}  // namespace heuristic
}  // namespace m1une

#endif  // M1UNE_HEURISTIC_BEAM_SEARCH_HPP
#line 1 "heuristic/beam_search.hpp"



#include <algorithm>
#include <cassert>
#include <concepts>
#include <cstddef>
#include <functional>
#include <type_traits>
#include <utility>
#include <vector>

#line 1 "heuristic/objective.hpp"



namespace m1une {
namespace heuristic {

enum class Objective {
    minimize,
    maximize,
};

template <class Score>
bool better_score(const Score& first, const Score& second,
                  Objective objective) {
    if (objective == Objective::maximize) return second < first;
    return first < second;
}

}  // namespace heuristic
}  // namespace m1une


#line 14 "heuristic/beam_search.hpp"

namespace m1une {
namespace heuristic {

template <class State, class Score>
struct BeamSearchResult {
    State state;
    Score score;
    int depth;
    std::size_t expanded_states;
    std::size_t generated_states;
};

namespace beam_search_detail {

template <class State, class Score>
struct Node {
    State state;
    Score score;
    std::size_t order;
};

template <class State, class Score>
struct BetterNode {
    Objective objective;

    bool operator()(const Node<State, Score>& first,
                    const Node<State, Score>& second) const {
        if (better_score(first.score, second.score, objective)) return true;
        if (better_score(second.score, first.score, objective)) return false;
        return first.order < second.order;
    }
};

}  // namespace beam_search_detail

// expand(state, next_depth) may return a range of children. For allocation-free
// generation, expand(state, next_depth, emit) may instead call emit(child).
// evaluate(state) returns its score. The best beam_width states are retained at
// every depth, and the best state in the last non-empty layer is returned.
template <class State, class Expand, class Evaluate>
auto beam_search(State initial_state, int depth_limit, int beam_width,
                 Expand expand, Evaluate evaluate,
                 Objective objective = Objective::maximize) {
    assert(0 <= depth_limit);
    assert(0 < beam_width);

    using Score = std::remove_cvref_t<
        std::invoke_result_t<Evaluate&, const State&>>;
    using Node = beam_search_detail::Node<State, Score>;
    using Better = beam_search_detail::BetterNode<State, Score>;

    Score initial_score = std::invoke(evaluate, initial_state);
    std::vector<Node> beam;
    beam.push_back(Node{std::move(initial_state),
                        std::move(initial_score), 0});

    std::size_t expanded_states = 0;
    std::size_t generated_states = 0;
    int reached_depth = 0;
    if (depth_limit < 0 || beam_width <= 0) depth_limit = 0;

    Better better{objective};
    for (int next_depth = 1; next_depth <= depth_limit; next_depth++) {
        std::vector<Node> candidates;
        candidates.reserve(static_cast<std::size_t>(beam_width));
        std::size_t order = 0;

        for (const Node& node : beam) {
            expanded_states++;
            auto emit = [&](auto&& candidate_state) {
                using Candidate = decltype(candidate_state);
                static_assert(std::is_constructible_v<State, Candidate>);
                State state(std::forward<Candidate>(candidate_state));
                Score candidate_score = std::invoke(evaluate, state);
                Node candidate{std::move(state), std::move(candidate_score),
                               order++};
                generated_states++;
                if (int(candidates.size()) < beam_width) {
                    candidates.push_back(std::move(candidate));
                    std::push_heap(candidates.begin(), candidates.end(), better);
                } else if (better(candidate, candidates.front())) {
                    std::pop_heap(candidates.begin(), candidates.end(), better);
                    candidates.back() = std::move(candidate);
                    std::push_heap(candidates.begin(), candidates.end(), better);
                }
            };
            if constexpr (std::invocable<Expand&, const State&, int>) {
                auto next_states =
                    std::invoke(expand, node.state, next_depth);
                for (auto& candidate_state : next_states) {
                    emit(std::move(candidate_state));
                }
            } else if constexpr (std::invocable<Expand&, const State&>) {
                auto next_states = std::invoke(expand, node.state);
                for (auto& candidate_state : next_states) {
                    emit(std::move(candidate_state));
                }
            } else {
                std::invoke(expand, node.state, next_depth, emit);
            }
        }

        if (candidates.empty()) break;
        beam = std::move(candidates);
        reached_depth = next_depth;
    }

    int best = 0;
    for (int index = 1; index < int(beam.size()); index++) {
        if (better(beam[index], beam[best])) best = index;
    }
    return BeamSearchResult<State, Score>{
        std::move(beam[best].state), std::move(beam[best].score),
        reached_depth, expanded_states, generated_states};
}

}  // namespace heuristic
}  // namespace m1une
Back to top page