approach.txt

approach:
it takes 0 jobs to print 0.
for each attendee count from 1 up to n:
  for each square not larger than attendee count:
    we can do it in 1 + jobs[attendee_count - square]
  take the smallest of the above possibilities
time:
proportional to n * sqrt(n)

kafka.cr

def kafka(n)
  prints = Array(UInt8).new(n + 1) { UInt8::MAX }
  prints[0] = 0

  isqrt = 1
  isqrt_increases = 2 * 2

  (1..n).each { |i|
    if isqrt_increases == i
      isqrt += 1
      isqrt_increases = (isqrt + 1) * (isqrt + 1)
    end

    best = prints[i - 1]
    (2..isqrt).each { |j|
      best = {best, prints[i - j * j]}.min
    }
    prints[i] = best + 1
  }
  prints
end

def prime_factors(n, primes)
  freq = Hash(Int32, Int32).new(0)

  primes.each { |prime|
    break if prime * prime > n
    while n % prime == 0
      freq[prime] += 1
      n /= prime
    end
  }

  freq[n] += 1 if n > 1
  freq
end

def primes_up_to(n)
  prime = Array.new(n + 1, true)
  prime[0] = false
  prime[1] = false

  2.step(to: Math.sqrt(n)) { |i|
    next unless prime[i]
    (i * i).step(to: n, by: i) { |j| prime[j] = false }
  }

  (0..n).select { |n| prime[n] }
end

n = ARGV.empty? ? 1_000_000 : ARGV[0].to_i32
primes = primes_up_to(n)

t = Time.now
kafka = kafka(n)
puts "generate #{n} in #{Time.now - t}"

t = Time.now
kafka.each_with_index { |prints, i|
  if i == 0
    puts "WRONG 0 #{prints}" if prints != 0
    next
  end

  puts "WRONG #{i} 0" if prints == 0

  decomp = prime_factors(i, primes)

  # obvious!
  square = decomp.all? { |_, power| power.even? }
  puts "WRONG #{i}! #{prints}" if (prints == 1) != square

  # https://en.wikipedia.org/wiki/Sum_of_two_squares_theorem
  should_two_squares = decomp.none? { |prime, power| power.odd? && prime % 4 == 3 }
  puts "WRONG #{i}!! #{prints}" if (prints <= 2) != should_two_squares

  # https://en.wikipedia.org/wiki/Legendre%27s_three-square_theorem
  should_three_squares = begin
    x = i
    while x % 4 == 0
      x /= 4
    end
    x % 8 != 7
  end
  puts "WRONG #{i}!!! #{prints}" if (prints <= 3) != should_three_squares

  # https://en.wikipedia.org/wiki/Lagrange%27s_four-square_theorem
  puts "WRONG #{i}!!!! #{prints}" if prints >= 5
}
puts "verify #{n} in #{Time.now - t}"

kafka.hs

import Control.Arrow ((&&&))
import Control.Monad (forM, forM_, when)
import Data.Array.IArray (assocs)
import Data.Array.MArray (newArray_, newListArray, readArray, writeArray)
import Data.Array.ST (runSTUArray)
import Data.List (group)
import System.Environment (getArgs)

kafka :: Int -> [(Int, Int)]
kafka n = assocs prints
  where prints = runSTUArray $ do
          a <- newArray_ (0, n)
          writeArray a 0 0
          forM_ [1 .. n] $ \i -> do
            candidates <- forM (smallerSquaresThan i) (\j -> readArray a (i - j * j))
            writeArray a i (1 + minimum candidates)
          return a
        smallerSquaresThan i = takeWhile (\k -> k * k <= i) [1..]

isqrt :: Int -> Int
isqrt = floor' . sqrt . fromIntegral
  -- This is just to avoid a compiler warning
  -- about a defaulted constraint
  where floor' = floor :: (Double -> Int)

primesUpTo :: Int -> [Int]
primesUpTo n = if n < 2 then [] else 2 : oddPrimes
  where oddPrimes = map fst $ filter snd $ assocs arr
        arr = runSTUArray $ do
          a <- newListArray (0, n) $ cycle [False, True]
          writeArray a 1 False
          forM_ [3, 5 .. isqrt n] $ \k -> do
            isPrime <- readArray a k
            when isPrime $ forM_ [k * k, k * k + k .. n] $ \mult ->
              writeArray a mult False
          return a

-- https://stackoverflow.com/questions/21276844/prime-factors-in-haskell
primeFactors :: [Int] -> Int -> [Int]
primeFactors [] _ = error "need more primes"
primeFactors (p:_) n | n < p * p = [n | n > 1]
primeFactors ps@(p:ps') n =
  let (q, r) = n `divMod` p
  in if r == 0 then p : primeFactors ps q else primeFactors ps' n

rm4 :: Int -> Int
rm4 n = let (q, r) = n `divMod` 4 in if r == 0 then rm4 q else n

main :: IO ()
main = do
  args <- getArgs
  let n = case args of
        [] -> 1000000
        a:_ -> read a
  let primes = primesUpTo n
  let (z1, z2):ks = kafka n
  when (z1 /= 0 || z2 /= 0) (putStrLn ("WRONG " ++ show (z1, z2)))
  forM_ ks $ \k@(i, prints) -> do
    when (prints == 0) (putStrLn ("WRONG " ++ show k))

    let decomp = (map (head &&& length) . group) (primeFactors primes i)

    -- obvious!
    let square = all (\(_, power) -> power `mod` 2 == 0) decomp
    when ((prints == 1) /= square) (putStrLn ("WRONG! " ++ show k))

    -- https://en.wikipedia.org/wiki/Sum_of_two_squares_theorem
    let shouldTwoSquares = all (\(prime, power) -> power `mod` 2 == 0 || prime `mod` 4 /= 3) decomp
    when ((prints <= 2) /= shouldTwoSquares) (putStrLn ("WRONG!! " ++ show k))

    -- https://en.wikipedia.org/wiki/Legendre%27s_three-square_theorem
    let shouldThreeSquares = (rm4 i) `mod` 8 /= 7
    when ((prints <= 3) /= shouldThreeSquares) (putStrLn ("WRONG!!! " ++ show k))

    -- https://en.wikipedia.org/wiki/Lagrange%27s_four-square_theorem
    when (prints >= 5) (putStrLn ("WRONG!!!! " ++ show k))

kafka.rb

# approach:
# it takes 0 jobs to print 0.
# for each attendee count from 1 up to n:
#   for each square not larger than attendee count:
#     we can do it in 1 + jobs[attendee_count - square]
#   take the smallest of the above possibilities
#
# time:
# proportional to n * sqrt(n)
def kafka(max)
  prints = Array.new(max + 1)
  #prevs = Array.new(max + 1)
  prints[0] = 0

  (1..max).each { |i|
    prints[i] = prints[i - 1] + 1
    #prevs[i] = i - 1

    (2..isqrt(i)).each { |j|
      new_prints = prints[i - j * j] + 1
      if new_prints < prints[i]
        prints[i] = new_prints
        #prevs[i] = i - j * j
      end
    }
  }
  #[prints, prevs]
  prints
end

def isqrt(n)
  (n ** 0.5).floor
end

n = ARGV.empty? ? 100_000 : Integer(ARGV[0])

t = Time.now
kafka = kafka(n)
puts "generate #{n} in #{Time.now - t}"

require 'prime'

t = Time.now
kafka.each_with_index { |prints, i|
  if i == 0
    puts "WRONG 0 #{prints}" if prints != 0
    next
  end

  puts "WRONG #{i} 0" if prints == 0

  decomp = Prime.prime_division(i)

  # obvious!
  square = decomp.all? { |_, power| power.even? }
  puts "WRONG #{i}! #{prints}" if (prints == 1) != square

  # https://en.wikipedia.org/wiki/Sum_of_two_squares_theorem
  should_two_squares = decomp.none? { |prime, power| power.odd? && prime % 4 == 3 }
  puts "WRONG #{i}!! #{prints}" if (prints <= 2) != should_two_squares

  # https://en.wikipedia.org/wiki/Legendre%27s_three-square_theorem
  should_three_squares = begin
    x = i
    x /= 4 while x % 4 == 0
    x % 8 != 7
  end
  puts "WRONG #{i}!!! #{prints}" if (prints <= 3) != should_three_squares

  # https://en.wikipedia.org/wiki/Lagrange%27s_four-square_theorem
  puts "WRONG #{i}!!!! #{prints}" if prints >= 5
}
puts "verify #{n} in #{Time.now - t}"