Advertisement
Guest User

my Lua code which fails some tests

a guest
Jan 4th, 2017
147
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Lua 3.20 KB | None | 0 0
  1.  
  2. local sqrt = math.sqrt
  3.  
  4. --  array of primes
  5. local primes = setmetatable({2, 3, 5}, {__index =
  6.    function(t, k)
  7.       local p = t[#t]
  8.       for k = #t + 1, k do
  9.          repeat
  10.             p = p + 4 * p % 6
  11.             local i, root, is_compound = 2, sqrt(p) - 1.5
  12.             repeat
  13.                i = i + 1
  14.                local pp = t[i]
  15.                is_compound = p % pp == 0
  16.             until is_compound or root < pp
  17.          until not is_compound
  18.          t[k] = p
  19.       end
  20.       return p
  21.    end
  22. })
  23. local _ = primes[100000]   -- precalculate all p_k (it takes 1.2 seconds)
  24.  
  25.  
  26. local function Z_inv(value) -- -- 0 <= value < 1000000007
  27.    -- returns value^(-1) mod 1000000007
  28.    -- returns nothing if value = 0
  29.    local b, s, p, n = 1000000007, value, 0, 1
  30.    while s > 1 do
  31.       local r = b % s
  32.       p, n = n, p + (r - b) / s * n
  33.       b, s = s, r
  34.    end
  35.    if s == 1 then
  36.       return n % 1000000007
  37.    end
  38. end
  39.  
  40. local function Z_mul(A, B)  -- 0 <= A, B < 1000000007
  41.    local R = B % 65536
  42.    return ((B - R) / 65536 * A % 1000000007 * 65536 + A * R) % 1000000007
  43. end
  44.  
  45. local G = setmetatable({[0] = 1}, {__index =
  46.    function(t, k)
  47.       local g = t[#t]
  48.       for k = #t + 1, k do
  49.          local ii = Z_inv(primes[k]-1)
  50.          g = Z_mul(g, Z_mul(ii, ii))
  51.          t[k] = g
  52.       end
  53.       return g
  54.    end
  55. })
  56. local _ = G[100000]   --  precalculate all G[m] (it takes only 0.15 sec)
  57.  
  58. --~ io.input("input.txt")
  59.  
  60. -------------------------------------------------------------------------------------
  61. --  MAIN LOOP
  62. -------------------------------------------------------------------------------------
  63. -- We have 10+ seconds remaining for the following loop
  64.  
  65. for _ = 1, io.read"*n" do
  66.    local m, a = io.read("*n", "*n")
  67.    local result = G[m]       -- this value has already been precalculated
  68.    for k = 1, m do
  69.       ------------------------------------------------------------------------
  70.       -- The body of this inner loop is executed up to 3.9 million times
  71.       -- I need to increase its speed by 11% to pass all test cases
  72.       -- But I don't know how optimize it more
  73.       ------------------------------------------------------------------------
  74.       local p = primes[k]    -- this value has already been precalculated
  75.       local ak2 = a + k + 2  -- 3 <= ak2 < 2^18
  76.       -- Next 16 lines calculate:  W = p^ak2
  77.       local exponent = ak2
  78.       local mask = 131072    -- 2^17
  79.       while mask > exponent do
  80.          mask = mask / 2
  81.       end
  82.       exponent = exponent - mask
  83.       local W = p
  84.       while mask > 1 do
  85.          local R = W % 65536
  86.          W = ((W - R) / 65536 * W % 1000000007 * 65536 + W * R) % 1000000007
  87.          mask = mask / 2
  88.          if exponent >= mask then
  89.             exponent = exponent - mask
  90.             W = W * p % 1000000007
  91.          end
  92.       end
  93.       -- Next 3 lines calculate:  result = result * (W - 1 - (p - 1) * ak2)
  94.       local W = (W - 1 - (p - 1) * ak2) % 1000000007
  95.       local R = W % 65536
  96.       result = ((W - R) / 65536 * result % 1000000007 * 65536 + result * R) % 1000000007
  97.       ------------------------------------------------------------------------
  98.    end
  99.    print(result)
  100. end
  101.  
  102. --~ print(os.clock())
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement