Inner-Product Search
(data-structure/inner_product_search.hpp)

• View this file on GitHub
• Last update: 2024-01-08 01:32:22+09:00
• Include: #include "data-structure/inner_product_search.hpp"

Depends on

• Convex Hull Trick (Binary Search Tree) (data-structure/cht/convex_hull_trick_binsearchtree.hpp)

Code

#pragma once
#include <algorithm>
#include <cassert>
#include <limits>

#include "cht/convex_hull_trick_binsearchtree.hpp"

/**
 * @brief Inner-Product Search
 */
template <typename T>
class InnerProductSearch {
   public:
    void add(T a, T b) {
        ++n;
        a_min = std::min(a_min, a);
        a_max = std::max(a_max, a);
        cht_pos.add(a, b);
        cht_neg.add(-a, -b);
    }

    T query(T x, T y) const {
        assert(n > 0);
        if (y == 0) {
            return (x > 0 ? a_min : a_max) * x;
        } else if (y > 0) {
            return cht_pos.get(x / y) * y;
        } else {
            return -cht_neg.get(x / y) * y;
        }
    }

    int size() const { return n; }

   private:
    static constexpr T INF = std::numeric_limits<T>::max();

    int n = 0;
    T a_min = INF, a_max = -INF;
    ConvexHullTrick<T> cht_0, cht_pos, cht_neg;
};

#line 2 "data-structure/inner_product_search.hpp"
#include <algorithm>
#include <cassert>
#include <limits>

#line 4 "data-structure/cht/convex_hull_trick_binsearchtree.hpp"
#include <set>
#include <utility>

template <typename T>
class ConvexHullTrick {
   public:
    void add(T a, T b) {
        a = -a, b = -b;
        auto m = lines.insert({a, b, 0});
        auto l = m, r = m;
        ++r;
        while (update(m, r)) {
            r = lines.erase(r);
        }
        if (l != lines.begin() && update(--l, m)) {
            m = lines.erase(m);
            update(l, m);
        }
        m = l;
        while (l != lines.begin() && (--l)->p >= m->p) {
            update(l, lines.erase(m));
            m = l;
        }
    }

    T get(T x) const {
        assert(!lines.empty());
        auto it = *lines.lower_bound(x);
        return -(it.a * x + it.b);
    }

   private:
    struct Line {
        mutable T a, b;    // ax + b
        mutable double p;  // intersection point with the next line
        bool operator<(const Line& o) const { return a < o.a; }
        bool operator<(T x) const { return p < x; }
    };

    using iterator = typename std::multiset<Line, std::less<>>::iterator;
    static constexpr double INF = std::numeric_limits<double>::max() / 2;

    std::multiset<Line, std::less<>> lines;

    bool update(iterator x, iterator y) const {
        if (y == lines.end()) {
            x->p = INF;
            return false;
        }
        if (x->a == y->a) {
            x->p = (x->b > y->b ? INF : -INF);
        } else {
            x->p = 1.0 * (y->b - x->b) / (x->a - y->a);
        }
        return x->p >= y->p;
    }
};
#line 7 "data-structure/inner_product_search.hpp"

/**
 * @brief Inner-Product Search
 */
template <typename T>
class InnerProductSearch {
   public:
    void add(T a, T b) {
        ++n;
        a_min = std::min(a_min, a);
        a_max = std::max(a_max, a);
        cht_pos.add(a, b);
        cht_neg.add(-a, -b);
    }

    T query(T x, T y) const {
        assert(n > 0);
        if (y == 0) {
            return (x > 0 ? a_min : a_max) * x;
        } else if (y > 0) {
            return cht_pos.get(x / y) * y;
        } else {
            return -cht_neg.get(x / y) * y;
        }
    }

    int size() const { return n; }

   private:
    static constexpr T INF = std::numeric_limits<T>::max();

    int n = 0;
    T a_min = INF, a_max = -INF;
    ConvexHullTrick<T> cht_0, cht_pos, cht_neg;
};

Back to top page