disc_partition.rb

require 'priority_queue'

def solve(n, max_discs, l, t, m)
  dp = Array.new(n + 1) { a = Array.new(n + 1, Float::INFINITY); a[0] = 0; a }
  cut = Array.new(n + 1) { Array.new(n + 1) }
  cumulate_ts = t.inject([0]) { |result, tvalue| result << result.last + tvalue }
  (1..max_discs).each do |i|
    k = 1
    pq = PriorityQueue.new
    (1..n).each do |j|
      cost = dp[i - 1][j - 1]
      cost += 1 if j > 1 && m[j - 2] == m[j - 1]
      pq.push [cost, j], cost * n + (n - j) # cost が異なるなら cost が小さい方を、cost が等しいなら index が大きい方を返す
      while cumulate_ts[j] - cumulate_ts[k - 1] > l
        k += 1
      end

      while pq.min_key[1] < k
        pq.delete_min
      end

      min = pq.min_key
      if dp[i][j] > min[0]
        dp[i][j] = min[0]
        cut[i][j] = min[1]
      end
    end
    if dp[i][n] < Float::INFINITY # 解を見つけた
      partitions = []
      k = n
      i.downto(1).each do |j|
        partitions << [cut[j][k]-1..k-1] # cut の値や k は1始まりであることに注意
        k = cut[j][k] - 1
      end
      return [dp[i][n], partitions.reverse]
    end
  end
  nil
end

n = 5; max_discs = 5; l = 30
t = [10, 10, 10, 10, 10]
m = [1, 1, 2, 2, 2]

cost, sep = solve(n, max_discs, l, t, m)
p "cost: #{cost}, separation: #{sep}"