Ruby と Python で解く AtCoder JSC2019 Qual B 等差数列の和 逆元

はじめに

AtCoder Problems の Recommendation を利用して、過去の問題を解いています。
AtCoder さん、AtCoder Problems さん、ありがとうございます。

今回のお題

AtCoder JSC2019 Qual B - Kleene Inversion
Difficulty: 797

今回のテーマ、等差数列の和 + 逆元

Ruby

転倒数って初めて聞きますが、例1 3 2について、どのような転倒数になるのか見てみます。

K	3	2	1	合計
132			1	1
132 132		4	1	5
132 132 132	7	4	1	12

1 3 2が一つで1、二つ並んで1 3 2 1 3 2は4 + 1、三つならんで1 3 2 1 3 2 1 3 2は7 + 4 + 1と、等差数列の和になっていることが分かります。
等差数列の和の公式は、＿こちら＿ を参照しました。

ruby.rb

n, k = gets.split.map(&:to_i)
a = gets.split.map(&:to_i)
b = a + a
c = 0
(a.size - 1).times do |i|
  (i + 1).upto(a.size - 1) do |j|
    c += 1 if a[i] > a[j]
  end
end
d = 0
(b.size - 1).times do |i|
  (i + 1).upto(b.size - 1) do |j|
    d += 1 if b[i] > b[j]
  end
end
d -= c * 2
def modpow(n, p, mod)
  r = 1
  while p > 0
    r = r * n % mod if p & 1 == 1;
    n = n * n % mod
    p >>= 1
  end
  r
end
def modinv(n, mod)
  modpow(n, mod - 2, mod)
end
MOD = 1_000_000_007
ans = (k - 1) * d % MOD
ans += 2 * c
ans %= MOD
ans *= k
ans %= MOD
ans *= modinv(2, MOD)
ans %= MOD
puts ans

cal.rb

c = 0
(a.size - 1).times do |i|
  (i + 1).upto(a.size - 1) do |j|
    c += 1 if a[i] > a[j]
  end
end
d = 0
(b.size - 1).times do |i|
  (i + 1).upto(b.size - 1) do |j|
    d += 1 if b[i] > b[j]
  end
end
d -= c * 2

初項cと公差dはここで求めています。

wa.rb

ans = (k - 1) * d % MOD
ans += 2 * c
ans %= MOD
ans *= k
ans %= MOD
ans *= modinv(2, MOD)
ans %= MOD

等差数列の和の公式部分ですが、2による割り算があるため、逆元を求める必要があります。

Python

python.py

n, k = map(int, input().split())
a = list(map(int, input().split()))
b = a + a
c = 0
for i in range(len(a)):
    for j in range(i + 1, len(a)):
        if a[i] > a[j]:
            c += 1
d = 0
for i in range(len(b)):
    for j in range(i + 1, len(b)):
        if b[i] > b[j]:
            d += 1
d -= c * 2
def modpow(n, p, mod):
    r = 1
    while p > 0:
        if p & 1 == 1:
            r = r * n % mod
        n = n * n % mod
        p >>= 1
    return r
def modinv(n, mod):
    return modpow(n, mod - 2, mod)
MOD = 10 ** 9 + 7
ans = (k - 1) * d % MOD
ans += 2 * c
ans %= MOD
ans *= k
ans %= MOD
ans *= modinv(2, MOD)
ans %= MOD
print(ans)

PyPy3 (2.4.0)でも、そのまま使用できました。

	Ruby	Python	PyPy3
コード長 (Byte)	621	676	676
実行時間 (ms)	948	1903	322
メモリ (KB)	1916	3188	41692

まとめ

• JSC2019 Qual B を解いた
• Ruby に詳しくなった
• Python に詳しくなった

参照したサイト
４ 等差数列の和