そのループ中断条件、二重ループになってない?

4 min read読了の目安(約4400字

そんな事やる奴お前だけだろって話ですね。
今回解いていた問題はこれ。

https://atcoder.jp/contests/abc065/tasks/arc076_b

アホ丸出しでクラスカル法で書いてTLE。
よく見ると入力規則のせいで辺の数がどえらい事になっていたので辺の数を絞り、深さ優先探索で連結確認してTLE。
調べてみると、どうも通常の木よりも早いUnion Find木という物があるらしい。
つまりこの問題はクラスカル法、辺の削減、Union Findの複合問題だったわけですね。
Atcoder Problemsでもギリギリとはいえ青問題。灰色コーダーの手に負える問題ではありませんでした。

突き当たった問題

https://akhtikd.com/posts/ruby-de-union-find/

https://hai3.net/blog/ruby-union-find-tree/

とにかくカンニングでも何でもやりながら書き上げたのが以下です。

# frozen_string_literal: true

module UnionFind
  def has_all_united
    self.one?(&:negative?)
  end

  def unite(start_city, goal_city)
    # 両方根なら深い方が左に来る
    high_rank_city, low_rank_city = [root_of(start_city), root_of(goal_city)].sort
    self[high_rank_city] += self[low_rank_city]
    self[low_rank_city] = high_rank_city
  end

  def same_union?(city_a, city_b)
    root_of(city_a) == root_of(city_b)
  end

  private

    def root?(now_node)
      self[now_node].negative?
    end

    def root_of(now_node)
      return now_node if root?(now_node)

      self[now_node] = root_of(self[now_node])
    end
end

CITIES_COUNT = gets.to_i
CITIES_PLACE = Array.new(CITIES_COUNT) { gets.split.map(&:to_i) }.map.with_index {|grid, index| grid << index }
X_SORTED_CITIES = CITIES_PLACE.sort_by {|city| city[0] }
Y_SORTED_CITIES = CITIES_PLACE.sort_by {|city| city[1] }
edges = []
city_i = 0
while city_i <= CITIES_COUNT - 2
  edges << [X_SORTED_CITIES[city_i][2], X_SORTED_CITIES[city_i + 1][2], X_SORTED_CITIES[city_i + 1][0] - X_SORTED_CITIES[city_i][0]]
  edges << [Y_SORTED_CITIES[city_i][2], Y_SORTED_CITIES[city_i + 1][2], Y_SORTED_CITIES[city_i + 1][1] - Y_SORTED_CITIES[city_i][1]]
  city_i += 1
end

edges.sort_by! {|edge| edge[2] }

union_list = Array.new(CITIES_COUNT) { - 1 }
union_list.extend UnionFind

distance_sum = 0
until union_list.has_all_united
  shortest_start_city, shortest_goal_city, shortest_edge_length = edges.shift
  next if union_list.same_union?(shortest_start_city, shortest_goal_city)

  distance_sum += shortest_edge_length
  union_list.unite(shortest_start_city, shortest_goal_city)
end
puts distance_sum

定数が多くて見づらいとか、cityとnodeが混ざってて見づらいとか、座標と街番号を配列に入れんなハッシュにしろハッシュにとか、何でrank方式じゃなくてsize方式なんだとかそういうツッコミは受け付けていません。
何が問題かと言うと、これでTLEしたという事です。

解決

なんだUnion Findって世間で言われているほど早くないのか? と2箇所ほどいじったらいきなりAC。
通ったコードはこちら。

# frozen_string_literal: true

module UnionFind
  def unite(start_city, goal_city)
    high_rank_city = root_of(start_city)
    low_rank_city =  root_of(goal_city)
    high_rank_city, low_rank_city = low_rank_city, high_rank_city if low_rank_city < high_rank_city
    self[high_rank_city] += self[low_rank_city]
    self[low_rank_city] = high_rank_city
  end

  def same_union?(city_a, city_b)
    root_of(city_a) == root_of(city_b)
  end

  private

    def root?(now_node)
      self[now_node].negative?
    end

    def root_of(now_node)
      return now_node if root?(now_node)

      self[now_node] = root_of(self[now_node])
    end
end

CITIES_COUNT = gets.to_i
CITIES_PLACE = Array.new(CITIES_COUNT) { gets.split.map(&:to_i) }.map.with_index {|grid, index| grid << index }
X_SORTED_CITIES = CITIES_PLACE.sort_by {|city| city[0] }
Y_SORTED_CITIES = CITIES_PLACE.sort_by {|city| city[1] }
edges = []
city_i = 0
while city_i <= CITIES_COUNT - 2
  edges << [X_SORTED_CITIES[city_i][2], X_SORTED_CITIES[city_i + 1][2], X_SORTED_CITIES[city_i + 1][0] - X_SORTED_CITIES[city_i][0]]
  edges << [Y_SORTED_CITIES[city_i][2], Y_SORTED_CITIES[city_i + 1][2], Y_SORTED_CITIES[city_i + 1][1] - Y_SORTED_CITIES[city_i][1]]
  city_i += 1
end

edges.sort_by! {|edge| edge[2] }

union_list = Array.new(CITIES_COUNT) { - 1 }
union_list.extend UnionFind

distance_sum = 0
until edges.empty?
  shortest_start_city, shortest_goal_city, shortest_edge_length = edges.shift
  next if union_list.same_union?(shortest_start_city, shortest_goal_city)

  distance_sum += shortest_edge_length
  union_list.unite(shortest_start_city, shortest_goal_city)
end
puts distance_sum

変えた場所は2箇所だけ。

  • 木同士をつなげるunionメソッドにて、配列オブジェクトを作らず不等号で処理。
  • 最後のループの終了条件を「全ノード連結したら終了」から「全部の辺を見終わるまで終了しない」に変更。

考察

解決した理由はおそらくループ終了条件でしょう。
改善前のhas_all_unitedメソッドは中身がone? です。いちいち配列を全部なめて、根が複数ないか確認していました。実質的に二重ループです。
改善後の条件式はempty。なので配列の最初の一個を見るだけで終了します。
その分、「edges配列の先頭から辺を取り出し、両端のノードの親が同じ事を確認する」という計算をedges配列が空になるまで余計に繰り返す事になります。
が、ノードの根を比べる計算量は一回。edges配列の回数分だけ繰り返しても大したことはありません。(そもそもedges配列の長さは2N-2くらいでそんなに長くない。)
というわけで、「ループを中断していいか確認のために二重ループを回すのは本末転倒。ループを中断しない方が早い」という話でした。