Line data    Source code

       1             : #pragma once
       2             : 
       3             : #include <deque>
       4             : #include <ext/pb_ds/assoc_container.hpp>  // Common file
       5             : #include <ext/pb_ds/tree_policy.hpp>      // tree_order_statistics_node_update
       6             : #include <stdexcept>
       7             : 
       8             : /**
       9             :  * @brief Order-statistics trees.
      10             :  *
      11             :  * These data structures behave similarly to sets and maps, but they also
      12             :  * provide order statistics; i.e., you can ask for the k-th element in the set,
      13             :  * or the number of elements in the set less than x.
      14             :  *
      15             :  * The vast majority of logic required to implement order-statistics trees is
      16             :  * already hidden in a few GNU G++ headers, so we are just wrapping them here
      17             :  * into more convenient, STL-like abstractions.
      18             :  */
      19             : namespace utils::orderstat {
      20             : 
      21             : namespace gnu = __gnu_pbds;
      22             : 
      23             : /// @brief Internal definitions
      24             : namespace internal {
      25             : 
      26             : template<class Key, class T, class CMP>
      27             : using orderstat = gnu::tree<Key, T, CMP, gnu::rb_tree_tag, gnu::tree_order_statistics_node_update>;
      28             : 
      29             : template<class Key>
      30             : using orderstat_multiset = orderstat<Key, gnu::null_type, std::less_equal<Key>>;
      31             : 
      32             : }  // namespace internal
      33             : 
      34             : /// Useful definitions
      35             : 
      36             : /**
      37             :  * @brief Set with order statistics.
      38             :  */
      39             : template<class Key>
      40             : using set = internal::orderstat<Key, gnu::null_type, std::less<Key>>;
      41             : 
      42             : /**
      43             :  * @brief Multiset with order statistics.
      44             :  */
      45             : template<class Key>
      46             : struct multiset: public internal::orderstat_multiset<Key> {
      47             :     size_t erase(const Key& val) {
      48             :         auto   it  = this->lower_bound(val - 1);
      49             :         size_t ret = 0;
      50             :         while(*it == val) {
      51             :             internal::orderstat_multiset<Key>::erase(it);
      52             :             ++ret;
      53             :             it = this->lower_bound(val - 1);
      54             :         }
      55             :         return ret;
      56             :     }
      57             : };
      58             : 
      59             : /**
      60             :  * @brief Map with order statistics.
      61             :  */
      62             : template<class Key, class T, class Compare = std::less<Key>>
      63             : struct map: public internal::orderstat<Key, T, Compare> {
      64             :    private:
      65             :     typedef internal::orderstat<Key, T, Compare> parentClass;
      66             : 
      67             :    public:
      68             :     typedef typename parentClass::iterator iterator;
      69             : 
      70             :     const T& at(const Key& key) const {
      71             :         auto it = parentClass::find(key);
      72             :         if(it == parentClass::end())
      73             :             throw std::out_of_range("utils::orderstat::map::at");
      74             :         return it->second;
      75             :     }
      76             :     T& at(const Key& key) {
      77             :         auto it = parentClass::find(key);
      78             :         if(it == parentClass::end())
      79             :             throw std::out_of_range("utils::orderstat::map::at");
      80             :         return it->second;
      81             :     }
      82             : 
      83             :     template<class... Args>
      84             :     std::pair<iterator, bool> emplace(Args&&... args) {
      85             :         return parentClass::insert({std::forward<Args>(args)...});
      86             :     }
      87             : };
      88             : 
      89             : /**
      90             :  * @brief Multimap with order statistics.
      91             :  */
      92             : template<class Key, class T>
      93             : using multimap = internal::orderstat<Key, T, std::less_equal<Key>>;
      94             : 
      95             : // clang-format off
      96             : template<
      97             :     class T,
      98             :     class Compare = std::less<T>
      99             : >
     100             : // clang-format on
     101             : class queue {
     102             :     size_t        start = 0;
     103             :     std::deque<T> data;
     104             : 
     105             :     std::map<T, size_t, Compare> dataToCounter;
     106             : 
     107             :    public:
     108           0 :     void push(const T& val) {
     109           0 :         size_t i = start + data.size();
     110             : 
     111           0 :         if(dataToCounter.count(val))
     112           0 :             throw std::runtime_error("utils::orderstat::queue: duplicate value");
     113             : 
     114           0 :         dataToCounter[val] = i;
     115             : 
     116           0 :         data.emplace_back(val);
     117           0 :     }
     118             : 
     119             :     template<class... Args>
     120           0 :     void emplace(Args&&... args) {
     121           0 :         push(T(args...));
     122           0 :     }
     123             : 
     124           0 :     void pop() {
     125           0 :         assert(dataToCounter.erase(front()) == 1);
     126           0 :         data.pop_front();
     127           0 :         if(!data.empty()) {
     128           0 :             start = dataToCounter.at(data.front());
     129             :         } else {
     130           0 :             assert(dataToCounter.empty());
     131           0 :             start = 0;
     132             :         }
     133           0 :     }
     134             : 
     135           0 :     T& at(size_t i) {
     136           0 :         return data.at(i);
     137             :     }
     138           0 :     const T& at(size_t i) const {
     139           0 :         return data.at(i);
     140             :     }
     141             : 
     142             :     T&       operator[](size_t i) { return data[i]; }
     143             :     const T& operator[](size_t i) const { return data[i]; }
     144             : 
     145           0 :     size_t size() const { return data.size(); }
     146             : 
     147           0 :     bool empty() const { return size() == 0; }
     148             : 
     149           0 :     T&       front() { return at(0); }
     150           0 :     const T& front() const { return at(0); }
     151             : 
     152           0 :     size_t order_of(const T& val) const {
     153             :         size_t i;
     154             :         try {
     155           0 :             i = dataToCounter.at(val);
     156           0 :         } catch(const std::out_of_range& e) {
     157           0 :             throw std::out_of_range("utils::orderstat::queue: value not found");
     158             :         }
     159           0 :         return i - start;
     160             :     }
     161             : 
     162             :     // int erase(const T& val) {
     163             :     //     size_t i;
     164             :     //     try {
     165             :     //         i = order_of(val);
     166             :     //     } catch(const std::out_of_range& e) {
     167             :     //         return 0;
     168             :     //     }
     169             : 
     170             :     //     if(i == 0) {
     171             :     //         pop();
     172             :     //         return 1;
     173             :     //     }
     174             : 
     175             :     //     size_t pos = start + i;
     176             : 
     177             :     //     data.erase(data.begin() + i);
     178             : 
     179             :     //     assert(dataToCounter.erase(val) == 1);
     180             : 
     181             :     //     for(auto& [k, v]: dataToCounter)
     182             :     //         if(v > pos)
     183             :     //             --v;
     184             : 
     185             :     //     return 1;
     186             :     // }
     187             : };
     188             : 
     189             : }  // namespace utils::orderstat