ABC356 D問題（Masked Popcount）を解く

AtCoder が提供しているABC（AtCoder Beginner Contest）356 D問題をC++とPythonで解いてみました。ABC356は、2024年6月1日21:00に実施されました。

AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。

D問題 Keys（Difficulty : 886）

問題の詳細は、リンク先をご覧ください。

ABC356 D問題 Masked Popcount

2進数の並びから法則を見つけて解きました。AtCoder Problems による Difficulty は 886 でした。

解答案

$M$ は、すべてのビットが立っている十分大きな数とします。$N=18$ までの二進数表現を確認します。

それぞれの数字ごとに1となっているビットをカウントするのでは、とても間に合いません。それぞれの桁に注目します。表を列でみます。

下から2桁目（表では1桁目と表現）は、長さが4のパターンが繰り返されています。下から3桁目（表では2桁目と表現）は、長さが8のパターンが繰り返されています。

• 表で d 桁目の列について
  • 長さ $2^{(d+1)}$ のパターンが繰り返される。それぞれのパターンには、$2^d$ 個の1が含まれている。
  • 長さ $2^{(d+1)}$ のパターンの余りの部分は、前半が0、後半が1となる。

C++ プログラム例（ABC356D）

上の考察をプログラムに変換します。

• $M$ に含まれるビット位置を配列 tm に格納します（14ー19行目）。
• $M$ に含まれているビットに関して
  • t1 : $2^d$
  • t2 : 含まれるパターンの個数
  • t3 : パターンに含まれない余りの数
  • 変数 result を更新します（28ー29行目）。

以下が、C++プログラムです。

#include <bits/stdc++.h>
#include <atcoder/modint>
using namespace std;
using namespace atcoder;

typedef long long int ll;
typedef modint998244353 mint;

int main()
{
	ll n, m;
	cin >> n >> m;

	vector<int> tm;
	for (int i = 0; i <= 60; ++i) {
		if ((m & (1LL << i)) != 0) {
			tm.push_back(i);
		}
	}

	mint result = 0;
	for (int i = 0; i < tm.size(); ++i) {
		ll t1, t2, t3;
		t1 = 1LL << tm[i];
		t2 = (n + 1) / (t1 * 2);
		t3 = (n + 1) % (t1 * 2);
		
		result += t2 * t1;
		result += max(0LL, t3 - t1);
	}

	cout << result.val() << endl;

	return 0;
}

AC（Accepted＝正しいプログラム）と判定されました。

Python プログラム例（ABC356D）

Python版も基本的な考え方は同じです。以下がプログラムです。

"""AtCoder Beginner Contest 356 D"""
MOD = 998244353
n, m = map(int, input().split())

tm = []
for i in range(61):
    if (m & (1 << i)) != 0:
        tm.append(i)

result = 0
for e in tm:
    t1 = 1 << e
    t2 = (n + 1) // (t1 * 2)
    t3 = (n + 1) % (t1 * 2)
    result += t2 * t1
    result %= MOD
    result += max(0, t3 - t1)
    result %= MOD

print(result)

こちらも「AC」と判定されました。

最後に

この問題は、時間を使いましたがコンテスト中に解けませんでした。アイデアは浮かんでいましたが、プログラムとして実装することができませんでした。

解けなかったことを過度に気にせず、コンテストに取り組んでいきます。

引き続き ABC の問題を紹介していきます。