2591 -- Set Definition

http://acm.pku.edu.cn/JudgeOnline/problem?id=2591
次のように定められた集合 S の N 番目(昇順)の要素を出力する問題.

  • 1 は S の要素である
  • x が S の要素ならば,2x + 1 と 3x + 1 も S の要素である
  • その他は S の要素でない


N の上限が 10000000 なので,最初に 10000000 番目まで S の要素を作っておいて答えた.

#include <iostream>
#include <vector>
using namespace std;

int main(void)
{
  static const int N = 10000000;
  vector<int> v(N);
  int a = 0, b = 0;
  v[0] = 1;
  for (int i = 1; i < N; i++) {
    int x = min(2*v[a]+1, 3*v[b]+1);
    if (x == 2*v[a]+1) {
      a++;
    }
    if (x == 3*v[b]+1) {
      b++;
    }
    v[i] = x;
  }

  int n;
  while (cin >> n) {
    cout << v[n-1] << endl;
  }
  return 0;
}


ちなみに Haskell なら以下のようにもっとシンプルに定義通りに書ける (これを言いたかっただけ).

s = 1 : merge (map (\x -> 2*x+1) s) (map (\x -> 3*x+1) s)
  where
    merge a@(x:xs) b@(y:ys) =
      case compare x y of
        LT -> x : merge xs b
        EQ -> x : merge xs ys
        GT -> y : merge a ys

main = interact (unlines . map (show . (s !!) . subtract 1 . read) . lines)