はじめに
AtCoder Problems の Recommendation を利用して、過去の問題を解いています。
AtCoder さん、AtCoder Problems さん、ありがとうございます。
今回のお題
AtCoder Beginner Contest D - Rain Flows into Dams
Difficulty: 945
今回のテーマ、累積和
連立方程式を解く問題で、numpy
やscipy
のlupライブラリ
で求めることができます。
但し、競プロらしくそれではTLE
になりますので、連立方程式を次のように変形することにより解きます。
b1 + b2 = a1
b2 + b3 = a2
b3 + b1 = a3
b1 = a1 + -a2 + a3
b2 = a1 + a2 + -a3
b3 = -a1 + a2 + a3
b2 = a1 - a2 + a3
= 2 * a1 - (a1 + a2 - a3)
= 2 * a1 - b1
Ruby (AC)
ruby.rb
n = gets.to_i
a = gets.split.map(&:to_i)
ans = [0] * n
a.each_with_index do |x, i|
if i.even?
ans[0] += x
else
ans[0] -= x
end
end
1.upto(n - 1) do |i|
ans[i] = 2 * a[i - 1] - ans[i - 1]
end
puts ans * ' '
rui.rb
a.each_with_index do |x, i|
if i.even?
ans[0] += x
else
ans[0] -= x
end
end
インデックスによって、加算・減算する特殊な累積和です。
Ruby Matrix(TLE)
TLE.rb
require 'matrix'
n = gets.to_i
a = gets.split.map(&:to_i)
b = []
n.pred.times do |i|
b[i] = [0] * i + [1] * 2 + [0] * (n - i - 2)
end
b << [1] + [0] * (n - 2) + [1]
m = Matrix.rows(b).lup.solve(a)
puts (m * 2).to_a.map{|e| e.to_i}.join(' ')
Matrix
ライブラリのlup
を使用した解法ですが、TLE
です。
Ruby BigDecimal(TLE)
TLE.rb
require 'bigdecimal'
require 'bigdecimal/util'
require 'bigdecimal/ludcmp'
include LUSolve
n = gets.to_i
a = gets.chomp.split.map(&:to_d)
zero = '0.0'.to_d
one = '1.0'.to_d
b = []
n.pred.times do |i|
b[i] = [zero] * i + [one] * 2 + [zero] * (n - i - 2)
end
b << [one] + [zero] * (n - 2) + [one]
b.flatten!
x = lusolve(b, a, ludecomp(b, a.size, zero, one), zero)
puts x.map{|e| e.to_i * 2}.join(' ')
LUSolve
を使用した解法ですが、TLE
です。
Python (AC)
python.py
from sys import stdin
def main():
input = stdin.readline
n = int(input())
a = list(map(int, input().split()))
ans = [0] * n
for i in range(n):
if i % 2 == 0:
ans[0] += a[i]
else:
ans[0] -= a[i]
for i in range(1, n):
ans[i] = 2 * a[i - 1] - ans[i - 1]
print(' '.join(map(str, ans)))
main()
Python numpy(TLE)
TLE.py
from sys import stdin
def main():
import numpy as np
input = stdin.readline
n = int(input())
a = [[] for _ in range(n)]
for i in range(n - 1):
a[i] = [0] * i + [1, 1] + [0] * (n - i - 2)
a[n - 1] = [1] + [0] * (n - 2) + [1]
A = np.array(a)
b = np.array(list(map(int, input().split())))
x = np.linalg.solve(A, b) * 2
print(' '.join(map(str, map(int, x))))
main()
numpy
のlinalgライブラリ
を使用した解法ですが、TLE
です。
Python scipy(TLE)
TLE.py
from sys import stdin
def main():
import numpy as np
from scipy import linalg
input = stdin.readline
n = int(input())
a = [[] for _ in range(n)]
for i in range(n - 1):
a[i] = [0] * i + [1, 1] + [0] * (n - i - 2)
a[n - 1] = [1] + [0] * (n - 2) + [1]
A = np.array(a)
b = np.array(list(map(int, input().split())))
LU = linalg.lu_factor(A)
x = linalg.lu_solve(LU, b) * 2
print(' '.join(map(str, map(int, x))))
main()
scipy
のlinalgライブラリ
を使用した解法ですが、TLE
です。
Ruby | Python | |
---|---|---|
コード長 (Byte) | 234 | 382 |
実行時間 (ms) | 130 | 98 |
メモリ (KB) | 14084 | 19132 |
まとめ
- ABC 138 D を解いた
- Ruby に詳しくなった
- Python に詳しくなった
参照したサイト
instance method Matrix#lup
連立一次方程式と線形重回帰をPythonで理解する