Ruby と Python で解く AtCoder AISING2020 D 繰返し二乗法

はじめに

エイシング プログラミング コンテスト 2020 に参加しました。
AtCoder さん、AtCoder Problems さん、ありがとうございます。

今回のお題

AtCoder エイシング プログラミング コンテスト 2020 D - Anything Goes to Zero
Difficulty: 1294

今回のテーマ、繰返し二乗法

	0	1	1	0	1	10進数表記
元の数値	$$0*2^4$$	$$1*2^3$$	$$1*2^2$$	$$0*2^1$$	$$1*2^0$$	13
0桁目ビット反転	$$1*2^4$$					13+16=29
1桁目ビット反転		$$0*2^3$$				13- 8= 5
2桁目ビット反転			$$0*2^2$$			13- 4= 9
3桁目ビット反転				$$1*2^1$$		13+ 2=15
4桁目ビット反転					$$0*2^0$$	13- 1=12
与えられた数値を01101としますと、上の表の様に分解できます。
また、分配則として下記が成り立ちますので、
$$(a+b)%m=(a%m+b%m)%m$$
ビット反転した桁の数値を出し入れすることにより、高速に全体の数値を求めることができます。

Ruby

ruby.rb

n = gets.to_i
x = gets.chomp
xs = x.to_i(2)
xc = x.count('1')
def mpow(n, p, mod)
  return 0 if mod.zero?
  r = 1
  while p > 0
    r = r * n % mod if p & 1 == 1
    n = n * n % mod
    p >>= 1
  end
  r
end
def popcount(u)
  return 0 if u.zero?
  a = u % u.to_s(2).count('1')
  1 + popcount(a)
end
xsp = mpow(xs, 1, xc + 1)
xsm = mpow(xs, 1, xc - 1)
n.times do |i|
  if x[i] == '0'
    y = xsp + mpow(2, n - i - 1, xc + 1)
    y = mpow(y, 1, xc + 1)
  elsif xc == 1
    puts '0'
    next
  else
    y = xsm - mpow(2, n - i - 1, xc - 1)
    y = mpow(y, 1, xc - 1)
  end
  puts popcount(y) + 1
end

以前投稿しました Ruby と Perl と Java と Python で解く AtCoder ATC 002 B のmpowメソッドをコピペ使用します。

mpow.rb

def mpow(n, p, mod)
  return 0 if mod.zero?
  r = 1
  while p > 0
    r = r * n % mod if p & 1 == 1
    n = n * n % mod
    p >>= 1
  end
  r
end

但し、今回は除数が0になる可能性がありますので、それに対応しています。

inout.rb

n.times do |i|
  if x[i] == '0'
    y = xsp + mpow(2, n - i - 1, xc + 1)
    y = mpow(y, 1, xc + 1)
  elsif xc == 1
    puts '0'
    next
  else
    y = xsm - mpow(2, n - i - 1, xc - 1)
    y = mpow(y, 1, xc - 1)
  end
  puts popcount(y) + 1
end

ここで、ビット反転した桁の数値の出し入れを行っています。
ここでも、除数が0の場合の処理が必要です。

popcount.rb

def popcount(u)
  return 0 if u.zero?
  a = u % u.to_s(2).count('1')
  1 + popcount(a)
end

2回目以降のpopcountを再帰計算しています。
ここをメモ化するともう少し速くなります。

Python

python.py

from sys import stdin

def mpow(n, p, mod):
    if mod == 0:
        return 0
    r = 1
    while p > 0:
        if p & 1 == 1:
            r = r * n % mod
        n = n * n % mod
        p >>= 1
    return r
def popcount(u):
    if u == 0:
        return 0
    a = u % bin(u).count('1')
    return 1 + popcount(a)
def main():
    input = stdin.readline
    n = int(input())
    x = '0b' + input()
    xs = int(x, 0)
    xc = x.count('1')
    xsp = mpow(xs, 1, xc + 1)
    xsm = mpow(xs, 1, xc - 1)
    for i in range(2, n + 2):
        if x[i] == '0':
            y = xsp + mpow(2, n - i + 1, xc + 1)
            y = mpow(y, 1, xc + 1)
        elif xc == 1:
            print(0)
            continue
        else:
            y = xsm - mpow(2, n - i + 1, xc - 1)
            y = mpow(y, 1, xc - 1)
        print(popcount(y) + 1)
main()

	Ruby	Python
コード長 (Byte)	629	872
実行時間 (ms)	625	1048
メモリ (KB)	15392	9636

まとめ

• AISING2020 D を解いた
• Ruby に詳しくなった
• Python に詳しくなった

参照したサイト
pythonで2進数を表す