Library

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

View the Project on GitHub noshi91/Library

:heavy_check_mark: test/w_ary_tree_set.pred.test.cpp

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/predecessor_problem"

#include "data_structure/w_ary_tree_set.cpp"

#include "other/int_alias.cpp"
#include <algorithm>
#include <iostream>
#include <string>
#include <vector>

int main() {
#include "other/fast_ios.cpp"

  w_ary_tree_set<10000000> wats;

  usize N, Q;
  std::cin >> N >> Q;
  std::string T;
  std::cin >> T;
  for (usize i = 0; i != N; i += 1) {
    if (T[i] == '1') {
      wats.insert(i);
    }
  }
  for (usize i = 0; i != Q; i += 1) {
    u32 type;
    usize k;
    std::cin >> type >> k;
    if (type == 0) {
      wats.insert(k);
    } else if (type == 1) {
      wats.erase(k);
    } else if (type == 2) {
      std::cout << wats.contains(k) << "\n";
    } else if (type == 3) {
      std::cout << static_cast<isize>(wats.succ(k)) << "\n";
    } else {
      std::cout << static_cast<isize>(wats.pred(k + 1)) << "\n";
    }
  }

  return 0;
}
#line 1 "test/w_ary_tree_set.pred.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/predecessor_problem"

#line 1 "other/ceildiv.cpp"
template <class T> constexpr T ceildiv(const T &n, const T &d) {
  return n / d + (n % d != 0 ? 1 : 0);
}
#line 2 "other/int_alias.cpp"

#include <cstddef>
#include <cstdint>

using i32 = std::int32_t;
using i64 = std::int64_t;
using u32 = std::uint32_t;
using u64 = std::uint64_t;
using isize = std::ptrdiff_t;
using usize = std::size_t;
#line 3 "data_structure/w_ary_tree_set.cpp"

#include <array>
#include <type_traits>

namespace w_ary_tree_set_impl {
constexpr usize calc_c(usize n) {
  usize ret = 1;
  while (n > 64) {
    ret *= 64;
    n /= 64;
  }
  return ret;
}
usize bsf(const u64 x) {
#ifdef __GNUC__
  return __builtin_ctzll(x);
#else
  unsigned long ret;
  if (_BitScanForward(&ret, x & ~(~static_cast<u64>(0) << 32))) {
    return ret;
  } else {
    _BitScanForward(&ret, x >> 32);
    return ret + 32;
  }
#endif
}
usize bsr(const u64 x) {
#ifdef __GNUC__
  return 63 - __builtin_clzll(x);
#else
  unsigned long ret;
  if (_BitScanReverse(&ret, x >> 32)) {
    return ret + 32;
  } else {
    _BitScanReverse(&ret, x & ~(~static_cast<u64>(0) << 32));
    return ret;
  }
#endif
}

template <usize N, class D = void> class w_ary_tree_set {
  static constexpr usize C = calc_c(N);

  u64 map;
  std::array<w_ary_tree_set<C>, ceildiv(N, C)> child;

public:
  w_ary_tree_set() : map(0), child() {}

  bool insert(const usize key) {
    const usize pos = key / C;
    map |= static_cast<u64>(1) << pos;
    return child[pos].insert(key % C);
  }

  bool erase(const usize key) {
    const usize pos = key / C;
    const bool ret = child[pos].erase(key % C);
    if (child[pos]._get_map() == 0) {
      map &= ~(static_cast<u64>(1) << pos);
    }
    return ret;
  }

  bool contains(const usize key) const {
    return child[key / C].contains(key % C);
  }

  usize min() const {
    const usize pos = bsf(map);
    return pos * C + child[pos].min();
  }

  usize max() const {
    const usize pos = bsr(map);
    return pos * C + child[pos].max();
  }

  usize pred(const usize key) const {
    const usize pos = key / C;
    const usize t = child[pos].pred(key % C);
    if (t != -1) {
      return pos * C + t;
    }
    const u64 masked = map & ~(~static_cast<u64>(0) << pos);
    if (masked == 0) {
      return -1;
    }
    const usize pos2 = bsr(masked);
    return pos2 * C + child[pos2].max();
  }

  usize succ(const usize key) const {
    const usize pos = key / C;
    const usize t = child[pos].succ(key % C);
    if (t != -1) {
      return pos * C + t;
    }
    const u64 masked = map & ~(~static_cast<u64>(0) >> (63 - pos));
    if (masked == 0) {
      return -1;
    }
    const usize pos2 = bsf(masked);
    return pos2 * C + child[pos2].min();
  }

  u64 _get_map() const { return map; }
};

template <usize N>
class w_ary_tree_set<N, typename std::enable_if<(N <= 64)>::type> {
  u64 map;

public:
  w_ary_tree_set() : map(0) {}

  bool insert(const usize key) {
    const u64 pop = static_cast<u64>(1) << key;
    if ((map & pop) != 0) {
      return false;
    } else {
      map |= pop;
      return true;
    }
  }

  bool erase(const usize key) {
    const u64 pop = static_cast<u64>(1) << key;
    if ((map & pop) != 0) {
      map &= ~pop;
      return true;
    } else {
      return false;
    }
  }

  bool contains(const usize key) const {
    return (map & static_cast<u64>(1) << key) != 0;
  }

  usize min() const { return bsf(map); }

  usize max() const { return bsr(map); }

  usize pred(const usize key) const {
    const u64 masked = map & ~(~static_cast<u64>(0) << key);
    if (masked == 0) {
      return -1;
    }
    return bsr(masked);
  }

  usize succ(const usize key) const {
    const u64 masked = map & (~static_cast<u64>(0) << key);
    if (masked == 0) {
      return -1;
    }
    return bsf(masked);
  }

  u64 _get_map() const { return map; }
};

} // namespace w_ary_tree_set_impl


template <usize N>
using w_ary_tree_set = w_ary_tree_set_impl::w_ary_tree_set<N>;

/**
 * @brief w-ary Tree Set
 * @docs docs/w_ary_tree_set.md
 */
#line 4 "test/w_ary_tree_set.pred.test.cpp"

#line 6 "test/w_ary_tree_set.pred.test.cpp"
#include <algorithm>
#include <iostream>
#include <string>
#include <vector>

int main() {
#line 1 "other/fast_ios.cpp"
std::ios::sync_with_stdio(false);
std::cin.tie(nullptr);
#line 13 "test/w_ary_tree_set.pred.test.cpp"

  w_ary_tree_set<10000000> wats;

  usize N, Q;
  std::cin >> N >> Q;
  std::string T;
  std::cin >> T;
  for (usize i = 0; i != N; i += 1) {
    if (T[i] == '1') {
      wats.insert(i);
    }
  }
  for (usize i = 0; i != Q; i += 1) {
    u32 type;
    usize k;
    std::cin >> type >> k;
    if (type == 0) {
      wats.insert(k);
    } else if (type == 1) {
      wats.erase(k);
    } else if (type == 2) {
      std::cout << wats.contains(k) << "\n";
    } else if (type == 3) {
      std::cout << static_cast<isize>(wats.succ(k)) << "\n";
    } else {
      std::cout << static_cast<isize>(wats.pred(k + 1)) << "\n";
    }
  }

  return 0;
}
Back to top page