spec/dsu_spec.cr

• View this file on GitHub
• Last update: 2023-11-28 17:27:51+09:00

Depends on

• src/dsu.cr

Required by

• spec/test_runner.cr

Code

# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

require "../src/dsu.cr"
require "spec"

alias DSU = AtCoder::DSU

describe "DSU" do
  describe "#merge" do
    it "unifies group of trees into unions" do
      tree = DSU.new(9_i64)

      tree.merge(0, 2)
      tree.merge(2, 4)
      tree.merge(4, 6)

      tree.merge(7, 1)
      tree.merge(7, 3)
      tree.merge(7, 5)

      tree.same?(2, 6).should eq true
      tree.same?(4, 2).should eq true
      tree.same?(1, 3).should eq true
      tree.same?(3, 5).should eq true
      tree.same?(8, 8).should eq true
      tree.same?(2, 1).should eq false
      tree.same?(7, 0).should eq false
      tree.same?(5, 4).should eq false
      tree.same?(8, 7).should eq false

      tree.size(4).should eq 4
      tree.size(1).should eq 4
      tree.size(8).should eq 1

      tree.groups.should eq Set{
        Set{0, 2, 4, 6},
        Set{1, 3, 5, 7},
        Set{8},
      }
    end
  end
end

# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# require "../src/dsu.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

module AtCoder
  # Implements [atcoder::dsu](https://atcoder.github.io/ac-library/master/document_en/dsu.html).
  #
  # ```
  # dsu = AtCoder::DSU.new(10)
  # dsu.merge(0, 2)
  # dsu.merge(4, 2)
  # dsu.same?(0, 4) # => true
  # dsu.size(4)     # => 3
  # ```
  class DSU
    getter parents : Array(Int64)
    getter sizes : Array(Int64)
    getter size : Int64

    def initialize(@size)
      @parents = Array.new(size, &.to_i64)
      @sizes = Array.new(size, 1_i64)
    end

    # Implements atcoder::dsu.leader(node).
    def leader(node)
      until @parents[node] == node
        @parents[node] = @parents[@parents[node]]
        node = @parents[node]
      end
      node
    end

    # Implements atcoder::dsu.merge(a, b).
    def merge(a, b)
      leader_a = leader(a.to_i64)
      leader_b = leader(b.to_i64)
      unless leader_a == leader_b
        if @sizes[leader_a] < @sizes[leader_b]
          leader_a, leader_b = leader_b, leader_a
        end
        @parents[leader_b] = leader_a
        @sizes[leader_a] += @sizes[leader_b]
      end
    end

    # Implements atcoder::dsu.same(a, b).
    def same?(a, b)
      leader(a) == leader(b)
    end

    # Implements atcoder::dsu.size().
    def size(node)
      @sizes[leader(node)]
    end

    # Implements atcoder::dsu.groups().
    # This method returns `Set` instead of list.
    def groups
      groups = Hash(Int64, Set(Int64)).new { |h, k| h[k] = Set(Int64).new }
      @size.times do |i|
        groups[leader(i)] << i
      end
      groups.values.to_set
    end
  end
end

require "spec"

alias DSU = AtCoder::DSU

describe "DSU" do
  describe "#merge" do
    it "unifies group of trees into unions" do
      tree = DSU.new(9_i64)

      tree.merge(0, 2)
      tree.merge(2, 4)
      tree.merge(4, 6)

      tree.merge(7, 1)
      tree.merge(7, 3)
      tree.merge(7, 5)

      tree.same?(2, 6).should eq true
      tree.same?(4, 2).should eq true
      tree.same?(1, 3).should eq true
      tree.same?(3, 5).should eq true
      tree.same?(8, 8).should eq true
      tree.same?(2, 1).should eq false
      tree.same?(7, 0).should eq false
      tree.same?(5, 4).should eq false
      tree.same?(8, 7).should eq false

      tree.size(4).should eq 4
      tree.size(1).should eq 4
      tree.size(8).should eq 1

      tree.groups.should eq Set{
        Set{0, 2, 4, 6},
        Set{1, 3, 5, 7},
        Set{8},
      }
    end
  end
end

Back to top page