問題のおさらい

こんにちは。堀江貴文です。

今回は、エンジニアのみなさんにぜひ解いていただきたい問題をCodeIQと一緒に考えてみました。
ちょっと考え方に工夫が必要な問題になっていますので、ぜひ挑戦してみてください。

【3つの棒をつなぎ合わせる組み合わせを数えてください】
長さの異なった棒がいくつかあります。
たとえば、次のように長さが8, 4, 10, 3, 2の5種類の棒があるとすると、これらの中から3つ選んでつなぎ合わせ、長さがちょうど15になる棒の長さの組み合わせは(8, 4, 3)と(10, 3, 2)の2通りあります。

つなぎ合わせる際の、棒の順番は考慮しません。つまり、(8, 4, 3), (8, 3, 4), (4, 8, 3), ...のような選び方は、すべて同じ組み合わせと考えます。また、つなぎ合わせる棒は必ず3本で、1本・2本だけで目的の長さになっていたとしても、それは組み合わせに含みません。
 
【問題】
 
つなぎ合わせたときの長さLと、N個(1≦N≦5000)の棒の長さが標準入力から与えられるとき、N個の棒の中から3つをつなぎ合わせて長さがLになる組み合わせの総数を求めるプログラムを書いてください。ただし、個々の棒の長さや、つなぎ合わせた長さ(L)は正の整数で、32bit整数で十分扱える範囲です。また、棒の長さはすべて異なるものとします。
 
【入力】
 
標準入力から、以下の形式で複数の整数値を読み込みます。

1行目のLは、つなぎ合わせた棒の長さです。
 
2行目のNは、棒の数です。
 
3行目以降のN行は個々の棒の長さで、長さはすべて異なるものとします。
 
【出力】
 
標準出力に、整数値を1つ出力してください。行末での改行は有無を問いません。
 
【入出力サンプル】
Case 1: Input
L
N
a1
a2
‥
‥
aN

Case 1: Output
2


Case 2: Input
35
10
13
12
17
10
4
18
3
11
5
7
Case 2: Output
6

※(4, 13, 18), (5, 12, 18), (5, 13, 17), (7, 10, 18), (7, 11, 17), (10, 12, 13)の6通り
 
 
【解答方法】
・データが標準入力で与えらた際、結果を標準出力に出力するプログラムを作成してください。
　★テストデータは標準入力より与えられる想定で記述して下さい。
・挑戦言語は不問です（ただし下記のプログラミング言語選択で選択可能なものに限る）。
・自分の書いたプログラム言語を選択し、解答欄にソースコードを貼りつけたら、
　送信前に「提出前に確認」ボタンをクリックしてください。構文エラーがないかどうかチェックできます。
　「解答コードは正常に実行されました」というメッセージを確認の上、「解答を送信」ボタンで解答してください。
・この問題にはテストケースが2つ用意されています。2つすべてに通れば正解です！
・各言語の標準入出力の方法は、こちらを参考にしてください。
 
 
【注意】
◆本問は自動採点ですので、入出力フォーマットを守ってください
 
【入力】入力は標準入力から取るようにしてください。
【出力】正解となる数のみを出力してください。「こたえは…です」や「answer=」のような、不要な文字列が出力されていると不正解扱いになってしまいますのでご注意ください。
 
◆テストデータは非公開です
 
プログラムの実行テストで用いるデータは非公開ですので、入力フォーマットに従い任意のデータを読み込めるようになっていなければなりません。
 
◆制限時間にもご注意ください
 
プログラムの実行には時間制限が設けられています。
入力データのサイズが小さい場合は特に問題ありませんが、大き目のデータを扱う場合は、ある程度アルゴリズムに工夫が必要な場合があります。
 
 
【ヒント】
3つの棒を選ぶ組み合わせをすべて列挙しようとすると、計算の量が多くなり、時間制限にかかってしまいますので、少し工夫が必要です。
 
たとえば2つの棒を選ぶと、残りの棒の長さは（必要な長さ）－（2つの棒の長さの和）で計算できます。この長さの棒があるかどうかを調べるだけで数え上げることができますね。

一般的な解答例

「3つの棒の長さを足し合わせる組み合わせ」ということで、最初のステップとしては3重ループのコードを書くことが考えられます。しかし入力（棒の数）の最大値は5000で、単純に5000の3乗回（1000億回以上）も計算を行うと制限時間（1秒）内に解を求めることができません。

3つの棒の長さをそれぞれL1, L2, L3とし、その合計がLsumになるとき、「L3 = Lsum – L1 – L2」になります。これを利用すると、「長さL3の棒が存在するかどうかさえ分かればよい」ということになり、実際のコードは2重ループで済みます。

「長さL3の棒が存在するかどうか」はハッシュマップでも構いませんが、サイズLsumの配列を用いるとシンプルになります。

// C++14
// https://ideone.com/ycq94t
#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

int main() {
    int Lsum;
    int nSticks;
    vector<int> sticks;
    vector<int> exist(Lsum, 0); // 存在する場合は1、存在しない場合は0

    cin >> Lsum >> nSticks;
    for ( int i = 0; i < nSticks; ++i ) {
        int a; cin >> a;
        sticks.emplace_back( a );
        if ( a < Lsum ) exist[a] = 1;
    }
    sort ( sticks.begin(), sticks.end() );

    int count = 0;
    for ( int i = 0; i < nSticks-2; ++i ) {
        for ( int j = i+1; j < nSticks-1; ++j) {
            int rest = Lsum - sticks[i] - sticks[j];
            if ( rest < 1 ) continue;
            if ( rest > sticks[j] && exist[ rest ] == 1 ) ++count;
        }
    }
    cout << count << endl;

    return 0;
}

「exist」を用いて3本目の棒が存在するかどうかを調べていますが、そのままでは同じ棒を選んでしまう場合を許すことになります。

そこで事前に棒の長さをソートしておき、2つ目の棒の長さと残りの長さ(rest)を比較することで同じ棒を使わないようにしています。C++ではこの程度で正解が得られます。

スクリプト言語でも通したいよね！

挑戦者の皆さんの中には上記のようなアルゴリズムを実装したにも関わらず、時間制限にかかって不正解になってしまった方もいらっしゃるのではないでしょうか。試しに先程のアルゴリズムを、Rubyで実装してみます。

# Ruby
# https://ideone.com/EYJ2SA

s = STDIN.gets.to_i # Lsum
n = STDIN.gets.to_i # nSticks
v = Array.new # sticks
e = Array.new( s, false ) # exist

n.times{
  a = STDIN.gets.to_i
  v << a
  e[a] = true if a < s
}
v.sort!

c = 0 # count
(0..n-3).each{|i|
  (i+1..n-2).each{|j|
    r = s-v[i]-v[j]
    next if r < 1
    c+=1 if r > v[j] && e[r]
  }
}

p c

このコードは実行に4秒程度かかってしまい、不正解となってしまいます。

ある程度高速に実行できる言語に書き替えるのも1つの選択肢ですが、もう少し工夫して高速化にチャレンジしてみましょう。

アルゴリズムの改良

まず手始めに、ループの回数を少し減らしてみましょう。数え上げのループ部分に、1行加えます。

(0..n-3).each{|i|
  break if v[i..i+2].inject(:+) > s # 長さの合計がLsumを超えたら終わり
  (i+1..n-2).each{|j|
    r = s-v[i]-v[j]
    next if r < 1
    c+=1 if r > v[j] && e[r]
  }
}

「v[i..i+2].inject(:+)」は、v[i]+v[i+1]+v[i+2]のこと。

配列のi番目から隣り合った3要素を足し合わせることで、i番目以降の棒をつなぎ合わせてできる長さの最小値を求めることになります。

この値がLsumを超えてしまったら、それ以上調べても無意味なので即座にループを抜けるようにしておきます。この1行を加えるだけで1秒近く速くなりました。

もちろんこれだけでは不十分ですので、もう少しループ部分の計算回数を減らす工夫をしてみましょう。

(0..n-3).each{|i|
  a = i+1
  a+=1 while a < n-1 && v[i]+v[a]+v[n-1] < s

  b = n-2
  b-=1 while b >= a && v[i]+v[b]+v[b+1] > s

  (a..b).each{|j|
    c+=1 if e[s-v[i]-v[j]]
  }
}

内側のループ区間を、より厳密に決めます。棒の長さが昇順のとき、i番目の棒が短い場合、i番目に近い棒では長さが足りない場合が多く、無駄な計算を行う原因となります。

そこでa番目の棒について、v[i]+v[a]+v[n-1]（v[n-1]は配列の末尾で棒の中で最も長いもの）の値がs(Lsum)を超えない限りaを増やすようにします。b番目の棒については、v[i]+v[b]+v[b+1]がsを超えない最大になるように決めます。超えてしまうと調べる意味がないからです。

このように書き替えても、実は実行速度はほとんど変わりません（むしろ少し遅くなります）。しかし、このように書くことで内側のループ内に余計な条件式を書く必要がなくなり非常にシンプルになりました。

内側のループ内が最も多く実行される箇所ですから、これは非常に重要です。かわりにa, bを求める部分が複雑になってしまいましたので、この部分を改良します。

# Ruby
# https://ideone.com/hqX8PX

s = STDIN.gets.to_i # Lsum
n = STDIN.gets.to_i # nSticks
v = Array.new # sticks
e = Array.new( s, false ) # exist

n.times{
  a = STDIN.gets.to_i
  v << a
  e[a] = true if a < s
}
v.sort!

def search(l, r, f)
  while r-l > 1
    k = (l+r)/2
    f.call(k) ? (l=k) : (r=k) 
  end
  l
end

c = 0 # count
(0..n-3).each{|i|
  a = search(i+1, n-1, lambda{|k| v[i]+v[k]+v[n-1] < s})
  b = search(i+1, n-1, lambda{|k| v[i]+v[k]+v[k+1] < s})

  (a..b).each{|j|
    c+=1 if e[s-v[i]-v[j]]
  }
}

p c

a, bを求める部分を二分探索するコードに書き替えることで格段に速くなりました。これでRubyでも正解することができますが、念のためもう少し高速化しておきます。

# コード全体を実行結果は https://ideone.com/ckjVir
(0..n-3).each{|i|
  r = s-v[i]
  a = search(i+1, n-1, lambda{|k| v[k]+v[n-1] < r})
  b = search(i+1, n-1, lambda{|k| v[k]+v[k+1] < r})

  while a <= b
    c+=1 if e[r-v[a]]
    a+=1
  end
}

不要な引き算を減らすため、s-v[i]を保持しておきます。現状ではRubyのループ処理はwhile文が高速に動きますので、内側のループはより高速になるようwhile文に書きかえました。これで楽々1秒を切るようになりました。

他のスクリプト言語での解答例

同じアルゴリズムで、PythonやJavaScriptでも正解することができますので解答例をご紹介します。

# Python
# https://ideone.com/rATRAT
# Python
import bisect

s = int(input())
n = int(input())

e = [ False for i in range(0, s) ]
v = [ int(input()) for i in range(0, n) ]
for t in v:
  if t < s: e[t]=True

v = sorted(v)
z = [ a+b for (a, b) in zip(v[0:n-1], v[1:n]) ]

def solve(i):
  r = s-v[i]
  a = bisect.bisect_left(v, r-v[n-1], lo=i+1)
  b = bisect.bisect_right(z, r, lo=a)
  return [ e[r-w] for w in v[a:b]].count(True)

print(sum([ solve(i) for i in range(0, n-2)]))

Pythonではループ処理に内包表記を用いると非常に高速になります。二分探索にはbisectモジュールを用いるとシンプルかつ高速です。v[k]+v[k+1]に相当する計算は事前に計算しておく(z)と、bisect.bisect_rightの部分がスッキリします。

// JavaScript (spidermonkey)
// https://ideone.com/ih3gks
var s = readline()|0;
var n = readline()|0;
var e = new Array(s);
for(var i = 0; i < s; ++i) e[i] = false;
var v = new Array(n);
for(var i = 0; i < n; ++i){
  v[i] = readline()|0;
  if( v[i] < s ) e[ v[i] ] = true;
}
v.sort( function(a,b){ return(a-b); } );
print( solve(s, n, v, e) );

function search(l, r, f){
  while( r-l > 1 ){
    var k = (l+r)>>1;
    if( f(k) ) l = k;
    else r = k;
  }
  return l;
}

function solve(s, n, v, e) {
  var count=0;
  for(var i=0; i<n-2; ++i){
    var a = search(i+1, n-1, (function(k){return v[i]+v[k]+v[n-1] < s;}));
    var b = search(i+1, n-1, (function(k){return v[i]+v[k]+v[k+1] < s;}));

    var r = s-v[i];
    for(var j = a; j <= b; ++j)
      if( e[r-v[j]] ) ++count;
  }
  return count;
}

アルゴリズムを工夫し、言語に応じた細かい高速化テクニックを使えるようにしましょう。速くなると思いこんで変わったコードを書くより、素直に書いて最適化がうまくかかってくれた方が速いこともあるでしょう。

言語ごとに、同じ言語であっても処理系のバージョンごとに最適化の方法は異なり、将来的にも変化していくだろうと考えて、その時どきで最適化のかかりやすい実装を考えてみるのも面白いかもしれませんね。

堀江さんのITエンジニアへのあふれる思い、コードへの関心が語られたインタビュー記事もぜひ読んでくださいね。

• 【堀江貴文氏インタビュー】一つところにとどまって満足するエンジニアを、僕は信用しない。

CodeIQ運営事務局・季節行事を大切にする会からの問題も大人気！こちらもぜひ♪

自分の書いたコードを誰かに評価されたいエンジニアは、けっこう多い？

ITエンジニアのための実務スキル評価サービス『CodeIQ』で出題されている「コード銀行」問題に挑戦すると、あなたのコードが評価されます。

評価(1)出題者からの評価　　⇒評価フィードバック例を見る

• 企業ではたらくという観点からあなたのコードをチェックします
• フィードバックされた観点をふまえてコードを書くと世の中の企業にとって「いいコード」が書けるようになります

評価(2)企業からの評価　　⇒評価フィードバック例を見る

• 「あなたと一緒にはたらきたい」という企業からスカウトが届きます
• あなたのコードが社会でどこまで通用するか、リアルな評価が得られます

興味を持った方はこちらからチャレンジを！