結果

問題 No.800 四平方定理
コンテスト
ユーザー 👑 obakyan
提出日時 2019-06-06 21:30:06
言語 Lua
(LuaJit 2.1.1734355927)
結果
WA  
実行時間 -
コード長 2,138 bytes
コンパイル時間 31 ms
コンパイル使用メモリ 6,692 KB
実行使用メモリ 19,336 KB
最終ジャッジ日時 2024-10-01 15:10:16
合計ジャッジ時間 22,838 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2 WA * 1
other AC * 11 WA * 18 TLE * 1
権限があれば一括ダウンロードができます

ソースコード

diff # 

local mfl = math.floor
local mmi = math.min
local function lower_bound(ary, x, max)
  local num = #ary
  if num == 0 then return 1 end
  if(x <= ary[1]) then return 1 end
  if(ary[max] < x) then return max + 1 end
  local min = 1
  while(1 < max - min) do
    local mid = mfl((min + max) / 2)
    if(ary[mid] < x) then
      min = mid
    else
      max = mid
    end
  end
  return max
end

local n, d = io.read("*n", "*n")
local sq = {}
local same_ev, same_odd = {}, {}
local evens = {}
local odds = {}
local same_evlen, same_oddlen = 0, 0
for i = 1, n do
  sq[i] = i * i
  if 2 * sq[i] < d + n * n then
    if i % 2 == 0 then
      same_evlen = same_evlen + 1
      same_ev[same_evlen] = 2 * sq[i]
    else
      same_oddlen = same_oddlen + 1
      same_odd[same_oddlen] = 2 * sq[i]
    end
  end
end
local evlen, odlen = 0, 0
for i = 1, n - 1 do
  for j = i + 2, n, 2 do
    if d + n * n <= sq[i] + sq[j] then break end
    evlen = evlen + 1
    evens[evlen] = sq[i] + sq[j]
  end
  for j = i + 1, n, 2 do
    if d + n * n <= sq[i] + sq[j] then break end
    odlen = odlen + 1
    odds[odlen] = sq[i] + sq[j]
  end
end
table.sort(evens)
table.sort(odds)
local cnt = 0
for w = 1, n do
  local w2d = d + sq[w]
  local curmax_odd, curmax_even = odlen, evlen
  for z = 1, n do
    local tgt = w2d - sq[z]
    if tgt <= 0 then break end
    if tgt % 4 == 0 then
      local a = lower_bound(same_ev, tgt, same_evlen)
      if a <= same_evlen and same_ev[a] == tgt then
        cnt = cnt + 1
      end
      a = lower_bound(evens, tgt, curmax_even)
      curmax_even = mmi(evlen, a)
      for k = a, evlen do
        if evens[k] == tgt then
          cnt = cnt + 2
        else break
        end
      end
    elseif tgt % 8 == 2 then
      local a = lower_bound(same_odd, tgt, same_oddlen)
      if a <= same_oddlen and same_odd[a] == tgt then
        cnt = cnt + 1
      end
    elseif tgt % 4 == 1 then
      local a = lower_bound(odds, tgt, curmax_odd)
      curmax_odd = mmi(odlen, a)
      for k = a, odlen do
        if odds[k] == tgt then
          cnt = cnt + 2
        else break
        end
      end
    end
  end
end
print(cnt)
0