#866765 (Lua) No.2290 UnUnion Find

提出ソース

結果

問題	No.2290 UnUnion Find
ユーザー	👑 obakyan
提出日時	2023-05-05 23:54:57
言語	Lua (LuaJit 2.1.1696795921)
結果	AC
実行時間	709 ms / 2,000 ms
コード長	11,771 bytes
コンパイル時間	91 ms
コンパイル使用メモリ	5,120 KB
実行使用メモリ	25,088 KB
最終ジャッジ日時	2024-05-02 19:06:40
合計ジャッジ時間	25,464 ms
ジャッジサーバーID （参考情報）	judge2 / judge3

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	1 ms 5,248 KB
testcase_01	AC	2 ms 5,376 KB
testcase_02	AC	111 ms 5,376 KB
testcase_03	AC	429 ms 14,160 KB
testcase_04	AC	435 ms 14,288 KB
testcase_05	AC	433 ms 14,260 KB
testcase_06	AC	430 ms 14,256 KB
testcase_07	AC	435 ms 14,248 KB
testcase_08	AC	430 ms 14,336 KB
testcase_09	AC	435 ms 14,288 KB
testcase_10	AC	436 ms 14,164 KB
testcase_11	AC	441 ms 14,256 KB
testcase_12	AC	445 ms 14,160 KB
testcase_13	AC	439 ms 14,380 KB
testcase_14	AC	450 ms 14,288 KB
testcase_15	AC	431 ms 14,256 KB
testcase_16	AC	439 ms 14,208 KB
testcase_17	AC	449 ms 14,384 KB
testcase_18	AC	437 ms 14,252 KB
testcase_19	AC	556 ms 14,508 KB
testcase_20	AC	708 ms 24,752 KB
testcase_21	AC	161 ms 5,376 KB
testcase_22	AC	344 ms 8,960 KB
testcase_23	AC	339 ms 9,088 KB
testcase_24	AC	628 ms 25,088 KB
testcase_25	AC	281 ms 8,832 KB
testcase_26	AC	667 ms 24,904 KB
testcase_27	AC	637 ms 24,720 KB
testcase_28	AC	464 ms 14,500 KB
testcase_29	AC	615 ms 24,880 KB
testcase_30	AC	192 ms 5,376 KB
testcase_31	AC	561 ms 14,420 KB
testcase_32	AC	299 ms 8,960 KB
testcase_33	AC	458 ms 14,512 KB
testcase_34	AC	709 ms 24,880 KB
testcase_35	AC	660 ms 24,784 KB
testcase_36	AC	310 ms 9,136 KB
testcase_37	AC	411 ms 14,592 KB
testcase_38	AC	306 ms 9,004 KB
testcase_39	AC	351 ms 8,960 KB
testcase_40	AC	643 ms 24,756 KB
testcase_41	AC	273 ms 8,876 KB
testcase_42	AC	529 ms 14,544 KB
testcase_43	AC	514 ms 14,716 KB
testcase_44	AC	433 ms 14,256 KB
testcase_45	AC	468 ms 14,124 KB
testcase_46	AC	492 ms 14,256 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

--[[
Reference:
"Red-Black Tree by Java & Python"
http://wwwa.pikara.ne.jp/okojisan/rb-tree/index.html
]]

local mfl, mce = math.floor, math.ceil
local band, bor = bit.band, bit.bor

local rb_bitmask_size = 30
local rb_bitmask = {1}
for i = 2, rb_bitmask_size do
  rb_bitmask[i] = rb_bitmask[i - 1] * 2
end

local RBSet = {}
RBSet.isRed = function(self, node_idx)
  local a = mce(node_idx / rb_bitmask_size)
  local b = node_idx - (a - 1) * rb_bitmask_size
  return 0 < band(self.colors[a], rb_bitmask[b])
end
RBSet.setRed = function(self, node_idx)
  local a = mce(node_idx / rb_bitmask_size)
  local b = node_idx - (a - 1) * rb_bitmask_size
  self.colors[a] = bit.bor(self.colors[a], rb_bitmask[b])
end
RBSet.setBlack = function(self, node_idx)
  local a = mce(node_idx / rb_bitmask_size)
  local b = node_idx - (a - 1) * rb_bitmask_size
  if 0 < band(self.colors[a], rb_bitmask[b]) then
    self.colors[a] = self.colors[a] - rb_bitmask[b]
  end
end
RBSet.rotateL = function(self, idx, right_idx)
  local l, r = self.l, self.r
  local k = self.key
  r[idx], r[right_idx], l[right_idx], l[idx] = r[right_idx], l[right_idx], l[idx], right_idx
  k[idx], k[right_idx] = k[right_idx], k[idx]
end
RBSet.rotateR = function(self, idx, left_idx)
  local l, r = self.l, self.r
  local k = self.key
  l[idx], l[left_idx], r[left_idx], r[idx] = l[left_idx], r[left_idx], r[idx], left_idx
  k[idx], k[left_idx] = k[left_idx], k[idx]
end
RBSet.rotateLR = function(self, idx, left_idx, leftright_idx)
  local l, r = self.l, self.r
  local k = self.key
  r[left_idx], l[leftright_idx], r[leftright_idx], r[idx] = l[leftright_idx], r[leftright_idx], r[idx], leftright_idx
  k[idx], k[leftright_idx] = k[leftright_idx], k[idx]
end
RBSet.rotateRL = function(self, idx, right_idx, rightleft_idx)
  local l, r = self.l, self.r
  local k = self.key
  l[right_idx], r[rightleft_idx], l[rightleft_idx], l[idx] = r[rightleft_idx], l[rightleft_idx], l[idx], rightleft_idx
  k[idx], k[rightleft_idx] = k[rightleft_idx], k[idx]
end
RBSet.create = function(self, lt)
  self.lt = lt
  if not lt then self.lt = function(a, b) return a < b end end
  self.node_count = 0
  self.capacity = 0
  self.root = 0
  self.free_nodes = {}
  self.l, self.r, self.key = {}, {}, {}
  self.colors = {}
end
RBSet.push = function(self, key)
  if self.node_count == self.capacity then
    self.capacity = self.capacity + 1
    if self.capacity % rb_bitmask_size == 1 then
      local z = mce(self.capacity / rb_bitmask_size)
      self.colors[z] = 0
    end
    table.insert(self.free_nodes, self.capacity)
  end
  local node_idxes_by_rank = {}
  local cur_idx = 0
  local cur_rank = 1
  local parent, granp = 0, 0
  local l, r, k = self.l, self.r, self.key
  local lt = self.lt
  cur_idx = self.root
  while 0 < cur_idx do
    node_idxes_by_rank[cur_rank] = cur_idx
    cur_rank = cur_rank + 1
    if lt(key, k[cur_idx]) then
      cur_idx = l[cur_idx]
    else
      cur_idx = r[cur_idx]
    end
  end
  self.node_count = self.node_count + 1
  cur_idx = self.free_nodes[self.node_count]
  node_idxes_by_rank[cur_rank] = cur_idx
  if cur_rank == 1 then
    self.root = cur_idx
  else
    parent = node_idxes_by_rank[cur_rank - 1]
    if lt(key, k[parent]) then
      l[parent] = cur_idx
    else
      r[parent] = cur_idx
    end
  end
  k[cur_idx], l[cur_idx], r[cur_idx] = key, 0, 0
  self:setRed(cur_idx)
  while true do
    cur_idx = node_idxes_by_rank[cur_rank]
    if cur_rank == 1 then
      self:setBlack(cur_idx)
      break
    end
    parent = node_idxes_by_rank[cur_rank - 1]
    if self:isRed(parent) then
      granp = node_idxes_by_rank[cur_rank - 2]
      if l[granp] == parent then
        if l[parent] == cur_idx then
          self:rotateR(granp, parent)
        else
          self:rotateLR(granp, parent, cur_idx)
        end
      else
        if l[parent] == cur_idx then
          self:rotateRL(granp, parent, cur_idx)
        else
          self:rotateL(granp, parent)
        end
      end
      self:setRed(granp)
      self:setBlack(parent)
      self:setBlack(cur_idx)
      cur_rank = cur_rank - 2
    else
      break
    end
  end
end

RBSet.remove = function(self, key)
  local l, r, k = self.l, self.r, self.key
  local lt = self.lt
  local node_idxes_by_rank = {}
  local cur_idx = 0
  local cur_rank = 1
  local parent, left, right, granc = 0, 0, 0, 0
  local resolve_state = 0 -- 0:OK, 1:left black count is small, 2:right is
  cur_idx = self.root
  while 0 < cur_idx do
    node_idxes_by_rank[cur_rank] = cur_idx
    if not lt(k[cur_idx], key) and not lt(key, k[cur_idx]) then
      break
    end
    cur_rank = cur_rank + 1
    if lt(key, k[cur_idx]) then
      cur_idx = l[cur_idx]
    else
      cur_idx = r[cur_idx]
    end
  end
  if 0 == cur_idx then
    return -- NOT FOUND
  end
  if 0 == l[cur_idx] then
    self.free_nodes[self.node_count] = cur_idx
    self.node_count = self.node_count - 1
    if cur_rank == 1 then
      self.root = r[cur_idx]
      if 0 < self.root then
        self:setBlack(self.root)
      end
      return
    else
      parent = node_idxes_by_rank[cur_rank - 1]
      if l[parent] == cur_idx then
        l[parent] = r[cur_idx]
        if not self:isRed(cur_idx) then
          resolve_state = 1
        end
      else
        r[parent] = r[cur_idx]
        if not self:isRed(cur_idx) then
          resolve_state = 2
        end
      end
      cur_idx = parent
      cur_rank = cur_rank - 1
    end
  else
    local swap_idx = cur_idx
    cur_idx = l[cur_idx]
    cur_rank = cur_rank + 1
    node_idxes_by_rank[cur_rank] = cur_idx
    while 0 < r[cur_idx] do
      cur_idx = r[cur_idx]
      cur_rank = cur_rank + 1
      node_idxes_by_rank[cur_rank] = cur_idx
    end
    k[swap_idx] = k[cur_idx]
    parent = node_idxes_by_rank[cur_rank - 1]
    if parent == swap_idx then
      l[parent] = l[cur_idx]
      if not self:isRed(cur_idx) then resolve_state = 1 end
    else
      r[parent] = l[cur_idx]
      if not self:isRed(cur_idx) then resolve_state = 2 end
    end
    self.free_nodes[self.node_count] = cur_idx
    self.node_count = self.node_count - 1
    cur_idx = parent
    cur_rank = cur_rank - 1
  end
  while 0 < resolve_state do
    if resolve_state == 1 then
      right = r[cur_idx]
      if self:isRed(right) then
        self:rotateL(cur_idx, right)
        cur_idx = right
        right = r[cur_idx]
        granc = l[right]
        if 0 < granc and self:isRed(granc) then
          self:rotateRL(cur_idx, right, granc)
          self:setBlack(granc)
          break
        end
        granc = r[right]
        if 0 < granc and self:isRed(granc) then
          self:rotateL(cur_idx, right)
          self:setBlack(granc)
          break
        end
        self:setBlack(cur_idx)
        self:setRed(right)
        break
      else
        granc = l[right]
        if 0 < granc and self:isRed(granc) then
          self:rotateRL(cur_idx, right, granc)
          self:setBlack(granc)
          break
        end
        granc = r[right]
        if 0 < granc and self:isRed(granc) then
          self:rotateL(cur_idx, right)
          self:setBlack(granc)
          break
        end
        self:setRed(right)
        if self:isRed(cur_idx) then
          self:setBlack(cur_idx)
          break
        end
        if cur_rank == 1 then break end
        cur_rank = cur_rank - 1
        parent = node_idxes_by_rank[cur_rank]
        if l[parent] == cur_idx then
          resolve_state = 1
        else
          resolve_state = 2
        end
        cur_idx = parent
      end
    else
      left = l[cur_idx]
      if self:isRed(left) then
        self:rotateR(cur_idx, left)
        cur_idx = left
        left = l[cur_idx]
        granc = r[left]
        if 0 < granc and self:isRed(granc) then
          self:rotateLR(cur_idx, left, granc)
          self:setBlack(granc)
          break
        end
        granc = l[left]
        if 0 < granc and self:isRed(granc) then
          self:rotateR(cur_idx, left)
          self:setBlack(granc)
          break
        end
        self:setBlack(cur_idx)
        self:setRed(left)
        break
      else
        granc = r[left]
        if 0 < granc and self:isRed(granc) then
          self:rotateLR(cur_idx, left, granc)
          self:setBlack(granc)
          break
        end
        granc = l[left]
        if 0 < granc and self:isRed(granc) then
          self:rotateR(cur_idx, left)
          self:setBlack(granc)
          break
        end
        self:setRed(left)
        if self:isRed(cur_idx) then
          self:setBlack(cur_idx)
          break
        end
        if cur_rank == 1 then break end
        cur_rank = cur_rank - 1
        parent = node_idxes_by_rank[cur_rank]
        if l[parent] == cur_idx then
          resolve_state = 1
        else
          resolve_state = 2
        end
        cur_idx = parent
      end
    end
  end
end

RBSet.getLeft = function(self, comp_key)
  local l, r, k = self.l, self.r, self.key
  local lt = self.lt
  local ret_key = nil
  local cur_idx = self.root
  while 0 < cur_idx do
    if not lt(comp_key, k[cur_idx]) then
      ret_key = k[cur_idx]
      cur_idx = r[cur_idx]
    else
      cur_idx = l[cur_idx]
    end
  end
  return ret_key
end

RBSet.getRight = function(self, comp_key)
  local l, r, k = self.l, self.r, self.key
  local lt = self.lt
  local ret_key = nil
  local cur_idx = self.root
  while 0 < cur_idx do
    if lt(k[cur_idx], comp_key) then
      cur_idx = r[cur_idx]
    else
      ret_key = k[cur_idx]
      cur_idx = l[cur_idx]
    end
  end
  return ret_key
end

RBSet.getMostLeft = function(self)
  local l, r, k = self.l, self.r, self.key
  local lt = self.lt
  local cur_idx = self.root
  local ret_key = nil
  while 0 < cur_idx do
    ret_key = k[cur_idx]
    cur_idx = l[cur_idx]
  end
  return ret_key
end

RBSet.getMostRight = function(self)
  local l, r, k = self.l, self.r, self.key
  local lt = self.lt
  local cur_idx = self.root
  local ret_key = nil
  while 0 < cur_idx do
    ret_key = k[cur_idx]
    cur_idx = r[cur_idx]
  end
  return ret_key
end

RBSet.enumerateDfs = function(self, p, t)
  if 0 < p then
    self:enumerateDfs(self.l[p], t)
    table.insert(t, self.key[p])
    self:enumerateDfs(self.r[p], t)
  end
end

RBSet.enumerateAll = function(self)
  local t = {}
  self:enumerateDfs(self.root, t)
  return t
end

local default_compare = function(a, b) return a < b end

RBSet.new = function(lt)
  if not lt then lt = default_compare end
  local obj = {}
  setmetatable(obj, {__index = RBSet})
  obj:create(lt)
  return obj
end

local UF = {}
UF.create = function(self, n)
  self.parent = {}
  self.rb = RBSet.new()
  for i = 1, n do
    self.parent[i] = i
    self.rb:push(i)
  end
  self.gc = n
end
UF.getroot = function(self, idx)
  local parent = self.parent
  local idx_update = idx
  while parent[idx] ~= idx do
    idx = parent[idx]
  end
  while parent[idx_update] ~= idx do
    parent[idx_update], idx_update = idx, parent[idx_update]
  end
  return idx
end
UF.unite = function(self, a, b)
  local ra = self:getroot(a)
  local rb = self:getroot(b)
  if ra ~= rb then
    self.rb:remove(rb)
    self.gc = self.gc - 1
    self.parent[b], self.parent[rb] = ra, ra
  end
end
UF.new = function(n)
  local obj = {}
  setmetatable(obj, {__index = UF})
  obj:create(n)
  return obj
end

local n = io.read("*n")
local uf = UF.new(n)
local q = io.read("*n")
for iq = 1, q do
  local tp = io.read("*n")
  if tp == 1 then
    local a, b = io.read("*n", "*n")
    uf:unite(a, b)
  else
    local a = io.read("*n")
    if 1 < uf.gc then
      local k1 = uf.rb:getMostLeft()
      local k2 = uf.rb:getMostRight()
      local r = uf:getroot(a)
      if r == k1 then
        print(k2)
      else
        print(k1)
      end
    else
      print(-1)
    end
  end
end

yukicoder

結果

テストケース

ソースコード