カメヲラボ

主にプログラミングとお勉強全般について書いてます

POJ3786: 隣接するビットを数えよう(ゴルフあり)

隣接するビットを数えよう

はじめに

この記事は、北京大学のオンラインでプログラミングの問題を解くサービスPKU JudgeOnline(POJ)の問題をピックアップし、解法の解説と、C/C++/Python等による実装例に加えて、少し短いCでの実装を紹介します。後半は完全に趣味の世界ですのでご注意ください。

本稿の構成

問題の説明

引用元

元はACM Regional Collegiate Programming Contest のF問題です。

問題内容

ビット列(0と1の並び)の中に、隣り合う1が何組あるかを答える問題です。入力はビット列の長さと、隣り合う1が何組あるかを表す数です。 例えば、"01101100"というビット列には、ビット列の長さが8で隣り合う1が2組あります。"0011100"なら、ビット列の長さは7で、隣り合う1の数は2と数えます。1が3つ並んでいる部分は、左2文字分で1組、右2文字分で1組の合計2組ということです。"1111"というビット列なら、隣り合う1は左2文字、中央2文字、右2文字の合計3組と数えます。つまり、1がx個連続していれば、隣り合う1の組みは(x-1)あるということです。

ビット列の長さ(n)は最大100で、隣り合うビットの組(k)はそれより少なくなります。条件としてnとkが与えられたとき、条件を満たすようなビットの並びは何通りあるかを計算して出力してください。

というような問題です。

入力例

以下のように、最初にテストケースの数が与えられ2行目以降に「テスト番号」「ビット列の長さ」「隣り合う1の組数」が空白文字で区切られています。

10 
1 5 2 
2 20 8 
3 30 17 
4 40 24 
5 50 37 
6 60 52 
7 70 59 
8 80 73 
9 90 84 
10 100 90

出力例

出力は、「テスト番号」「条件を満たすビットの並びの総数」を空白文字で区切り、各テストケースごとに改行します。

1 6
2 63426
3 1861225
4 168212501
5 44874764
6 160916
7 22937308
8 99167
9 15476
10 23076518

具体例

入力例の1つ目のテストケースでは、nが5でkが2です。このとき考えられるビットの並びは、 11100, 01110, 00111, 10111, 11101, 11011 となっているため、出力の1行目には6と出力されます。

アプローチ

長さnのビット列を全て列挙し、隣り合うビット列をカウントする方法は、自分で小さなテストケースを作成するのに良い方法ですが、それをそのまま解答とすることはできません。ビット列の長さは最大100もあるので、100ビット整数、つまり2100(=1267650600228229401496703205376)通りも調べなければならず、現実的ではありませんね。

DPによる解法と実装例

概要

隣り合う1の組を単純に数えるには非常に時間がかかりますし、テストケースも複数あります。一度数えたビット列に対して何度も数えなおさず済むように数を記録し、さらに長いビット列を調べる場合にそれを利用できるようにします。

表記方法

ビット列の長さをn、そのビット列に含まれる隣り合う1をk組とするときのビットの並びの総数をf(n, k)と表します。

準備

まず、長さ1のビット列を考えます。長さ1のビット列は、"0"と"1"の2通りありますが、隣り合うビット列は2文字以上並んでいないと数えられませんので、いずれの場合も隣り合う1は0組と考えます。

長さ2のビット列は、"00", "01", "10", "11"の4組ですが、この中で1が並んでいるのは"11"の1つだけで、それ以外の3つは0組です。この場合、f(2, 1) = 1, f(2, 0) = 3のように表すことができます。

長さ3まで見てみましょう。この場合は、"000", "001", "010", "011", "100", "101", "110", "111"の8通りで、f(3, 2) = 1, f(3, 1) = 2, f(3, 0) = 5となっています。

数え方を検討する

ビット列の長さが4, 5, 6, ...となると全部書き出すのは少し大変になってきますので、そろそろ数え方のルールを考えるときです。長さ4のビット列は全部で16通りありますが、それは長さ3のビット列8通りのそれぞれに対して、「先頭に0をつけたもの」と「先頭に1をつけたもの」の各8通りを合わせたものです。

このとき、先頭のビットに0をつけたものでは隣り合う1の組は増えませんが、先頭のビットに1をつけた場合は隣り合う1の組が1つ増える可能性があります。たとえば、"011"の先頭に1をつけたものは"1011"で隣り合う1の組は増えていませんが、"110"の先頭に1をつけた場合は"1110"となって、隣り合う1の組みが1から2に増えます。ここに着目すると、計算のルールが決められるはずです。

2種類のテーブルを用意する

ビット列の先頭に0を加えた場合は変化なしとして、1を加えた場合は隣り合う1の数が変化する可能性があるということを確認しました。変化するのは、1を加える前のビット列の最上位ビットが1の場合のみです。つまり、ビット列を次の2つのグループに分けて扱えば良さそうです。

  • 最上位ビットが0のビット列(dp0)

  • 最上位ビットが1のビット列(dp1)

これらについて、ビット列の長さnと隣り合う1の組数kをキーとするビット列の順列数を保持するテーブルを作成しておけば、f(n, k) = dp0[n][k] + dp1[n][k]のような計算で解が得られます。

各テーブルの更新方法

最上位ビットに0を加えたものは隣り合う1の組数が変わらないので、

dp0[n + 1][k] = dp0[n][k] + dp1[n][k]

で得られます。最上位ビットに1を加えたものは、

dp1[n + 1][k] = dp0[n][k] + dp1[n][k - 1]

のように計算します。dp1の添え字に注意してください。

実装例

以上の方法でPythonのコードを書いてみました。

import sys

# 問題がのサイズが最大100なのでそれに合わせた二次元配列
# dp[k][-1]で番兵にアクセスできるように実際のサイズはNより大きく
N = 100
SIZE = N + 2
dp0 = [[0] * SIZE for i in range(SIZE)]  # 最上位ビットが0のビット列用
dp1 = [[0] * SIZE for i in range(SIZE)]  # 最上位ビットが1のビット列用
dp0[1][0] = 1
dp1[1][0] = 1

def make_dp():
    for i in range(N):
        for j in range(i):
            dp0[i + 1][j] = dp0[i][j] + dp1[i][j]
            dp1[i + 1][j] = dp0[i][j] + dp1[i][j - 1]
        dp1[i + 1][i] = 1

def solve(n, k):
    return dp0[n][k] + dp1[n][k]

# テストケースの数
n_test = int(sys.stdin.readline());

make_dp()

for line in sys.stdin:
   tid, n, k = list(map(int, line.split()))
   print(f"{tid} {solve(n, k)}")

オリジナルの問題ではビット列の長さの上限が100ですので、それを定数にしています。また、Pythonではリストの添え字に-1を使うことでリストの最後尾にアクセスできることから、リストの最後尾に0を置いてループ内のコードがシンプルになるようにしています。

組み合わせ計算による解法と実装例

次は、任意のnとkに対して直接計算する方法を考えてみましょう。

はじめに

DPの場合と違って、ある程度大きなサイズのビット列で考えます。あまりに小さな例だと誤解する可能性があるからです。というわけで、まずはn=10, k=5くらいで考えてみましょう。この条件を満たすビット列の中で、ビット列に含まれる1の数が最も少ないのは"0001111110"のような、1が6つ連続し、他がすべて0であるような場合です。

少し一般的に言えば、長さnのビット列の中に含まれる1は少なくとも(k + 1)個あるということになります。では、ここから1の数を増やす操作を考えてみましょう。"0001111110"の"0"の部分を1に置き換えると、"1001111110"や"0101111110"は問題ありませんが、"0011111110"や"0001111111"の場合はkの値が変化してしまいます。kの値を増やさずにビット列の中の1を増やすには、どうすればよいでしょうか。

ビット列の分割(その1)

先の例を引き続き用いて、kの値、つまり隣り合う1の組数を変えずにビット列の1を増やすことを考えてみます。まずは1の並びが最小数6であるビット列"111111"に1を追加する方法を検討します。このビット列には、最上位ビットにも再開ビットにも1を加えるとkが増加してしまいますが、元のビット列のどこかに0を挿入すると、kの値は変化しません。例えば"111111"の中央に0を挿入してみると、"1110111"となります。このとき、隣う1の組は左側も右側もそれぞれ2になり、合計値すなわちkの値は2+2=4になります。さらにこのビット列の左右に1を加えると"11110111", "11101111"のようになり、いずれの場合もk=5をキープできます。

"0"を1つ挿入して左右のどちらかに"1"を付け加えるという操作であることは間違いないのですが、計算上は少しややこしいので、先に"1"を付け加えてから0を挿入すると考えます。そうすると、分割後に"1"を加える操作まで含めて6通りあるということになります。

今度は全体のビット長をnにするため"0"を補います。"0"を補っても良い箇所は、最初に"0"を挿入した部分と、左右両端の三か所です。

現時点でビット列の長さは8なので、あと2文字分の"0"を、左|中|右のどこかに挿入します。これは"00||"の並び方に相当し、C(4, 2) = 6という計算(Cは組み合わせの計算)で得られます。"11110111"に"0"を2つ挿入する場合、

0011110111
0111100111
0111101110
1111000111
1111001110
1111011100

のように、確かに6通りありました。 以上をまとめると、まず"1"の並びに"1"を1つ追加し、隙間に0を挿入する場合が6通りあって、それぞれの場合に残った"0"2つ分を挿入する場合が6通りあるので、全部で6*6=36通りあることになります。

ビット列の分割(その2)

では次に、元の"1"のビット列"111111"に、"1"を2つ付け加えることを考えてみます。先と同じように、"1"を2つ付け加えて"11111111"(1が8個)として0を2か所(同じ場所は×)に挿入し、"1"を並びを3分割してみます。このときの分け方は、C(7, 2) = 21通りあります。この場合は全てのビットが埋まりましたので、残りの"0"を挿入するステップが必要なくなります。"0"がないので"|||"の並び方と考えて計算しても、1通りという結果が得られます。

この考え方を使って、一般的な計算方法を考えてみましょう。 まず、長さnのビット列で隣り合う"1"の総数がkになるようなものの中で、"1"の数は(k+1)個が最小値です。このとき、m個の"1"を追加し、さらにm個の"0"を挿入して連続する"1"の並びを(m+1)個に分割します。

このとき、"1"の並びを分割する場合の数はC(k+m, m)であり、それぞれの場合に対して残りの"0"を挿入する場合の数はC(n-k-m, m+1)になります。したがって、

f(n, k) = C(k, 0) * C(n-k, 1) + C(k+1, 1) * C(n-k-1, 2) + C(k+2, 2) * C(n-k-2, 3) + ...

のように計算できます。

n=10, k=5の例なら、

f(10, 5)
= C(5, 0) * C(5, 1) + C(6, 1) * C(4, 2) + C(7, 2) * C(3, 3)
= 1 * 5 + 6 * 6 + 21 * 1
= 5 + 36 + 21
= 62

のようになります。

実装例

以上の方法でPythonのコードを書いてみました。

import sys
from scipy.special import comb

def solve(n, k):
    sum = 0
    a = n - k
    b = k
    c = 0
    while a >= c + 1:
        # exact=True で正確な整数値が得られる
        sum += comb(a, c + 1, exact=True) * comb(b, c, exact=True)
        a -= 1
        b += 1
        c += 1
    return sum

# テストケースの数
n_test = int(sys.stdin.readline());

for line in sys.stdin:
   tid, n, k = list(map(int, line.split()))
   print(f"{tid} {solve(n, k)}")

楽しい実装例

Cの実装

私の趣味はCで書かれたコードを短く書くことですので、実際に提出するコードをできるだけ短いCのコードに書き換えます。パッと見では組み合わせ計算の方が短く書けそうなので、これをまずCで書いてみます。

#include<stdio.h>

int combination(int n, int k) {
    long long int a = 1;
    for(int i=1; i<=k; ++i) {
        a = a * n / i;
        --n;
    }
    return a;
}

int solve(int n, int k) {
    int sum = 0;
    int a = n - k;
    int b = k;
    int c = 0;
    while(a >= c + 1) {
        sum += combination(a, c + 1) * combination(b, c);
        a -= 1;
        b += 1;
        c += 1;
    }
    return sum;
}

int main() {
    // テストケースの数
    int n_test;
    scanf("%d", &n_test);
    
    for(int i=0; i<n_test; ++i) {
        int tid, n, k;
        scanf("%d %d %d", &tid, &n, &k);
        printf("%d %d\n", tid, solve(n, k));
    }
    
    return 0;
}

長くなってしまいました(泣) というわけで、(POJ上で)ギリギリ動く範囲で省略できるものを省略していきます。

スッキリ書こう

f(n,k,m,j){
    for(m=j=1;j<=k;--n,++j)m=1.*m*n/j;
    return m;
}

g(n,k,s,c){
    n-=k;
    for(s=c=0;n>=c+1;--n,++k,++c)s+=f(n,c+1)*f(k,c);
    return s;
}

t,i,a,b;

main(z){
    for(;~scanf("%d",&b);a=b)--z%3||printf("%d %d\n"-!z,z/-3,g(a,b));
}

関数名をcombination->f, solve->gと書き換えました。それから変数の型を省略しました。全部int型です。とってもわかりやすいですね。しかし私は気に食わないんですよ。関数が3つもあることに!本当にひどいコードなので、main関数だけにしてしまいましょう。

関数っていらないよね

double x,j,u,w;

g(n,k,s,c){
    n-=k;
    for(s=c=0;n>=c+1;s+=x*(n-c)/j,--n,++k,++c)
        for(u=n,w=k,x=j=1;j<=c;++j,--u,--w)x=x*u/j*w/j;
    return s;
}

a,b;
main(z){for(;~scanf("%d",&b);--z%3||printf("%d %d\n"-!z,z/-3,g(a,b)),a=b);}

ひとまず関数gの中にfを埋め込みました。組み合わせの掛け算をまとめて行うための精度確保に泣く泣くdoubleを使いました。本当に悲しいですが、とりあえずgに消えてほしいです。

double x;
s,b,c,j,n,k,w;
main(z){
    for(;~scanf("%d",&b);--z%3?n=b:printf("%d %d\n"-!z,z/-3,s))
        for(k=b,s=c=0;x=j=n-k>c++;s+=x*c/w)
            for(w=++k;j+~c;)x=x*--w/j*(n-k-j+2)/j++;
}

だいぶスッキリしましたね!しかしまだまだ汚らしいコードに見えます。forって3回も書いてますし、doubleは本当に消えてほしい…。しかし現状私のスキルではこれくらいのサイズにしかできませんでした。空白文字や改行を取り除くと、このコードは166Bです。そしてこれより短いコードを書くことは可能です。よければ挑戦してみてください。

ところでコードゴルフってご存じですか?

あなたはコードゴルフをご存じですか?コードゴルフというのは、ソースコードをできるだけ短く書くという遊びです。「ソースコードを短く書く」⇒「少ない打鍵(ストローク)数でコードを書く」ということから、プログラミング上のゴルフと呼ばれています。

本稿の記事を読んで「ちょっと難しいなー」と思ったあなた!どうぞご安心ください。Anarchy Golf(通称あなごる)にはもっと簡単な問題もありますし、C以外の様々なプログラミング言語コードゴルフを楽しむこともできます。

興味を持った方は是非是非ご参加ください!