src/nglib/data_structure/dynamic_range_frequency.cr

View this file on GitHub
Last update: 2024-07-28 06:26:08+09:00

Depends on

src/nglib/data_structure/aatree_set.cr

Required by

src/nglib.cr

Code

require "./aatree_set.cr"

module NgLib
  # 長さ $n$ の整数列 $a_0, a_1, \cdots, a_{n-1}$ について、
  # $[l, r)$ に $x$ が何回現れるかを $O(\log{N})$ で計算するクラスです。
  class DynamicRangeFrequency(T)
    @size : Int32
    @map : Hash(Int32, NgLib::AATreeSet(Int32))
    @values : Array(T)
  
    def initialize(array : Array(T))
      @values = array.clone
      @size = array.size
      @map = Hash(Int32, NgLib::AATreeSet(Int32)).new
      array.each_with_index do |a, i|
        @map[a] = NgLib::AATreeSet(Int32).new unless @map.has_key?(a)
        @map[a] << i
      end
    end
  
    def count(range : Range(Int?, Int?), x : T)
      left = (range.begin || 0).to_i32
      right = ((range.end || @size) + (range.exclusive? || range.end.nil? ? 0 : 1)).to_i32
      v = @map[x]? || NgLib::AATreeSet(Int32).new
      lower_bound(v, right) - lower_bound(v, left)
    end
  
    def []=(i : Int, x : T)
      @map[@values[i]].delete(i.to_i32)
      @map[x] = NgLib::AATreeSet(Int32).new unless @map.has_key?(x)
      @map[x] << i.to_i32
      @values[i] = x
    end
  
    private def lower_bound(v : NgLib::AATreeSet(Int32), x : Int32)
      v.greater_equal_index(x) || v.size
    end
  end
end

# require "./aatree_set.cr"
module NgLib
  # 順序付き集合です。
  #
  # 平衡二分探索木として [AA木](https://ja.wikipedia.org/wiki/AA%E6%9C%A8) を使用しています。
  # 性能は赤黒木の方が良いことが多い気がします。
  #
  # C++の標準ライブラリの `multiset` と違って、$k$ 番目の値が取り出せることなどが魅力的です。
  class AATreeSet(T)
    include Enumerable(T)

    private class Node(T)
      property left : Node(T)?
      property right : Node(T)?
      property parent : Node(T)?
      property key : T
      property level : Int32
      property size : Int32

      def initialize(val : T)
        @left = @right = @parent = nil
        @level = 1
        @key = val
        @size = 1
      end

      def rotate_left : Node(T)
        right = @right.not_nil!
        mid = right.left
        par = @parent
        if right.parent = par
          if par.not_nil!.left == self
            par.not_nil!.left = right
          else
            par.not_nil!.right = right
          end
        end
        mid.parent = self if @right = mid
        right.left = self
        @parent = right

        sz = @size
        @size += (mid ? mid.size : 0) - right.size
        right.size = sz

        right
      end

      def rotate_right : Node(T)
        left = @left.not_nil!
        mid = left.right
        par = @parent

        if left.not_nil!.parent = par
          if par.not_nil!.left == self
            par.not_nil!.left = left
          else
            par.not_nil!.right = left
          end
        end
        mid.parent = self if @left = mid
        left.not_nil!.right = self
        @parent = left

        sz = @size
        @size += (mid ? mid.size : 0) - left.size
        left.size = sz

        left
      end

      def left_side?(node : Node(T)?) : Bool
        @left == node
      end

      def assign(node : Node(T)) : T
        @key = node.key
      end
    end

    @root : Node(T)?

    private def find_node(node : Node(T)?, val : T) : Node(T)?
      return nil unless node
      until val == node.not_nil!.key
        if val < node.not_nil!.key
          break unless node.not_nil!.left
          node = node.not_nil!.left
        else
          break unless node.not_nil!.right
          node = node.not_nil!.right
        end
      end
      node
    end

    private def skew(node : Node(T)?) : Node(T)?
      return nil unless node
      left = node.not_nil!.left
      if left && node.not_nil!.level == left.not_nil!.level
        return node.not_nil!.rotate_right
      end
      node
    end

    private def split(node : Node(T)?) : Node(T)?
      return nil unless node
      right = node.right
      if right && right.not_nil!.right && node.level == right.not_nil!.right.not_nil!.level
        r = node.rotate_left
        r.level += 1
        return r
      end
      node
    end

    private def begin_node : Node(T)?
      return nil unless @root
      node = @root
      while node.not_nil!.left
        node = node.not_nil!.left
      end
      node
    end

    private def next_node(node : Node(T)) : Node(T)?
      if node.right
        node = node.right
        while node.not_nil!.left
          node = node.not_nil!.left
        end
        node
      else
        while node
          par = node.not_nil!.parent
          if par && par.not_nil!.left_side?(node)
            return par
          end
          node = par
        end
        node
      end
    end

    private def level(node : Node(T)?)
      node ? node.level : 0
    end

    def initialize
      @root = nil
      self
    end

    def initialize(enumerable : Enumerable(T))
      @root = nil
      concat(enumerable)
      self
    end

    def concat(elems) : self
      elems.each { |elem| self << elem }
      self
    end

    def includes?(val : T) : Bool
      node = find_node(@root, val)
      node.nil? ? false : node.key == val
    end

    def clear
      @root = nil
    end

    def empty? : Bool
      @root.nil?
    end

    def at(k : Int) : T
      raise IndexError.new unless 0 <= k && k < size
      node = @root
      k += 1
      loop do
        left_size = (node.not_nil!.left ? node.not_nil!.left.not_nil!.size : 0) + 1
        break if left_size == k

        if k < left_size
          node = node.not_nil!.left
        else
          node = node.not_nil!.right
          k -= left_size
        end
      end
      node.not_nil!.key
    end

    def at?(k : Int) : T?
      return nil unless 0 <= k && k < size
      at(k)
    end

    def each(& : T ->)
      node = begin_node
      while node
        yield node.not_nil!.key
        node = next_node(node.not_nil!)
      end
    end

    def add(val : T) : Nil
      add?(val)
      nil
    end

    def add?(val : T) : Bool
      unless @root
        @root = Node.new(val)
        return true
      end

      node = find_node(@root, val)
      return false if node.not_nil!.key == val # NOT multi

      new_node = Node.new(val)
      if val <= node.not_nil!.key
        node.not_nil!.left = new_node
      else
        node.not_nil!.right = new_node
      end
      new_node.not_nil!.parent = node

      node = new_node
      while node
        node = split(skew(node))
        unless node.not_nil!.parent
          @root = node
          break
        end
        node = node.not_nil!.parent
        node.not_nil!.size += 1
      end
      true
    end

    def delete(val : T) : Bool
      return false unless @root

      node = find_node(@root, val)
      return false unless node.not_nil!.key == val

      if node.not_nil!.left || node.not_nil!.right
        child = find_node(node.not_nil!.left ? node.not_nil!.left : node.not_nil!.right, val)
        node.not_nil!.assign(child.not_nil!)
        node = child
      end

      par = node.not_nil!.parent
      if par
        if par.not_nil!.left_side?(node)
          par.left = nil
        else
          par.right = nil
        end
      else
        @root = nil
      end
      node = par

      while node
        new_level = {level(node.left), level(node.right)}.min + 1
        if new_level < node.level
          node.level = new_level
          if new_level < level(node.right)
            node.right.not_nil!.level = new_level
          end
        end

        node.size -= 1
        node = skew(node).not_nil!
        skew(node.right.not_nil!.right) if skew(node.right)

        node = split(node)
        split(node.not_nil!.right)

        unless node.not_nil!.parent
          @root = node
          break
        end
        node = node.not_nil!.parent
      end
      true
    end

    def delete_at(k : Int) : Bool
      delete(at(k))
    end

    def delete_at?(k : Int) : Bool
      val = at?(k)
      if val
        delete(val)
      else
        false
      end
    end

    def lower_bound_index(val : T) : Int32
      node = @root
      return 0 unless node
      index = 0
      while node
        if val <= node.not_nil!.key
          node = node.not_nil!.left
        else
          index += (node.not_nil!.left ? node.not_nil!.left.not_nil!.size : 0) + 1
          node = node.not_nil!.right
        end
      end
      index
    end

    def upper_bound_index(val : T) : Int32
      node = @root
      return 0 unless node
      index = 0
      while node
        if val < node.not_nil!.key
          node = node.not_nil!.left
        else
          index += (node.not_nil!.left ? node.not_nil!.left.not_nil!.size : 0) + 1
          node = node.not_nil!.right
        end
      end
      index
    end

    def less_index(val : T) : Int32?
      index = lower_bound_index(val)
      index == 0 ? nil : index - 1
    end

    def less_equal_index(val : T) : Int32?
      index = lower_bound_index(val)
      val == at?(index) ? index : (index == 0 ? nil : index - 1)
    end

    def greater_index(val : T) : Int32?
      index = upper_bound_index(val)
      index == size ? nil : index
    end

    def greater_equal_index(val : T) : Int32?
      index = lower_bound_index(val)
      index == size ? nil : index
    end

    def first : T
      at(0)
    end

    def first? : T?
      at?(0)
    end

    def last : T
      at(size - 1)
    end

    def last? : T?
      at?(size - 1)
    end

    def count(val : T) : Int32
      upper_bound_index(val) - lower_bound_index(val)
    end

    def size : Int32
      @root.try &.size || 0
    end

    def to_a : Array(T)
      res = Array(T).new
      return res unless @root
      dfs = uninitialized Node(T) -> Nil
      dfs = ->(node : Node(T)) do
        left = node.left
        right = node.right
        dfs.call(left) if left
        res << node.key
        dfs.call(right) if right
        nil
      end
      dfs.call(@root.not_nil!)
      res
    end

    def to_s(io : IO) : Nil
      io << "{" + to_a.join(", ") + "}"
    end

    def inspect(io : IO) : Nil
      to_s(io)
    end

    def ==(other : AATreeSet(T)) : Bool
      self.to_a == other.to_a
    end

    def <<(val : T) : Bool
      add?(val)
    end

    def [](k : Int) : T
      at(k)
    end

    def []?(k : Int) : T | Nil
      at?(k)
    end
  end
end

module NgLib
  # 長さ $n$ の整数列 $a_0, a_1, \cdots, a_{n-1}$ について、
  # $[l, r)$ に $x$ が何回現れるかを $O(\log{N})$ で計算するクラスです。
  class DynamicRangeFrequency(T)
    @size : Int32
    @map : Hash(Int32, NgLib::AATreeSet(Int32))
    @values : Array(T)
  
    def initialize(array : Array(T))
      @values = array.clone
      @size = array.size
      @map = Hash(Int32, NgLib::AATreeSet(Int32)).new
      array.each_with_index do |a, i|
        @map[a] = NgLib::AATreeSet(Int32).new unless @map.has_key?(a)
        @map[a] << i
      end
    end
  
    def count(range : Range(Int?, Int?), x : T)
      left = (range.begin || 0).to_i32
      right = ((range.end || @size) + (range.exclusive? || range.end.nil? ? 0 : 1)).to_i32
      v = @map[x]? || NgLib::AATreeSet(Int32).new
      lower_bound(v, right) - lower_bound(v, left)
    end
  
    def []=(i : Int, x : T)
      @map[@values[i]].delete(i.to_i32)
      @map[x] = NgLib::AATreeSet(Int32).new unless @map.has_key?(x)
      @map[x] << i.to_i32
      @values[i] = x
    end
  
    private def lower_bound(v : NgLib::AATreeSet(Int32), x : Int32)
      v.greater_equal_index(x) || v.size
    end
  end
end

ngng628's Library

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

src/nglib/data_structure/dynamic_range_frequency.cr

Depends on

Required by

Code