結果

問題 No.1307 Rotate and Accumulate
コンテスト
ユーザー 👑 obakyan
提出日時 2021-08-06 12:20:30
言語 Lua
(LuaJit 2.1.1734355927)
結果
AC  
実行時間 373 ms / 5,000 ms
コード長 2,724 bytes
コンパイル時間 413 ms
コンパイル使用メモリ 6,688 KB
実行使用メモリ 22,528 KB
最終ジャッジ日時 2024-09-17 00:28:18
合計ジャッジ時間 9,159 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 19
権限があれば一括ダウンロードができます

ソースコード

diff # 

local mfl = math.floor
local TFFT = {}

TFFT.initialize = function(self)
  self.size = 18
  -- 2^18
  self.n = 262144
  -- 1007681537: prime,
  -- 1007681537 % 262144 = 1
  self.mod = 1007681537
  -- (6161^262144) % mod = 1, (6161^131072) % mod ~= 1
  self.w = 6161
  -- (1007677693 * 262144) % mod = 1
  self.ninv = 1007677693
  -- (534835031 * 6161) % mod = 1
  self.winv = 534835031
  self.p2 = {1}
  for i = 2, self.size do
    self.p2[i] = self.p2[i - 1] * 2
  end
  self.binv = {}
  for i = 1, self.n do
    local y, z = 0, i - 1
    for j = 1, self.size do
      y = y + (z % 2) * self.p2[self.size + 1 - j]
      z = mfl(z / 2)
    end
    self.binv[i] = y + 1
  end
  self.wmul = {1}
  for i = 2, self.n do
    self.wmul[i] = self:mul(self.wmul[i - 1], self.w)
  end
  self.winvmul = {1}
  for i = 2, self.n do
    self.winvmul[i] = self:mul(self.winvmul[i - 1], self.winv)
  end
end

-- (44893^2) % 1007681537 = 18375
-- 1007681537 = 22446 * 44893 + rem(13259)
-- x0, y0 <= 44892
-- x1, y1 <= 22446
-- max(x1y1*18375+(x1y0+x0y1)*44893+x0y0) < 10^14
TFFT.mul = function(self, x, y)
  local x0, y0 = x % 44893, y % 44893
  local x1, y1 = mfl(x / 44893), mfl(y / 44893)
  return (x1 * y1 * 18375 + (x1 * y0 + x0 * y1) * 44893 + x0 * y0) % self.mod
end

TFFT.add = function(self, x, y) return (x + y) % self.mod end

TFFT.fft_common = function(self, ary, wmul)
  local ret = {}
  for i = 1, self.n do
    ret[i] = ary[self.binv[i]]
  end
  for i = 1, self.size do
    local step_size = self.p2[i]
    local step_count = self.p2[self.size + 1 - i]
    for istep = 1, step_count do
      local ofst = (istep - 1) * step_size * 2
      for j = 1, step_size do
        local a1, a2 = ret[ofst + j], ret[ofst + step_size + j]
        ret[ofst + j] = self:add(a1, self:mul(a2, wmul[1 + (j - 1) * step_count]))
        ret[ofst + step_size + j] = self:add(a1, self:mul(a2, wmul[1 + (j + step_size - 1) * step_count]))
      end
    end
  end
  return ret
end

TFFT.fft = function(self, ary)
  return self:fft_common(ary, self.wmul)
end
TFFT.ifft = function(self, ary)
  local ret = self:fft_common(ary, self.winvmul)
  for i = 1, self.n do
    ret[i] = self:mul(ret[i], self.ninv)
  end
  return ret
end

local n, q = io.read("*n", "*n")

TFFT:initialize()
local a, b = {}, {}
for i = 1, TFFT.n do
  a[i] = 0
  b[i] = 0
end
for i = 1, n do
  a[i] = io.read("*n")
end
for i = 1, q do
  local r = io.read("*n")
  if r == 0 then
    b[1] = b[1] + 1
  else
    b[n + 1 - r] = b[n + 1 - r] + 1
  end
end

local at = TFFT:fft(a)
local bt = TFFT:fft(b)
for i = 1, TFFT.n do
  at[i] = TFFT:mul(at[i], bt[i])
end
local c = TFFT:ifft(at)
for i = 1, n do
  c[i] = c[i] + c[i + n]
  io.write(c[i])
  io.write(i == n and "\n" or " ")
end
0