ICPC国内予選C その2

予定してませんでしたが。そして多分その2で終わりです。

昨日学内の練習会でImos Judge(http://judge.imoz.jp/)を利用させて頂いて国内予選の復習をしていたのですが、そこで「Cの制限時間が厳しくて通らない」という話題で少し盛り上がっていました。
厳しいといっても、他の問題の制限時間が10secの中でCだけ15secなので緩いはずなんですが、普通に書くと10! * 12 * 8になってしまって通りません。

しばらくするとチームメイトのH君と他チームコーチのT君が「アルファベットごとの係数を先に求めてしまえば良い」という事に気がつきました。これで10! * 10になって1/10位の時間で動いてくれるはずです（実際には12 * 8のテストケースはそれほど無いようで1/4程度にしかなりませんでしたが）。

例えば最初のテストケースだと

3
ACM
IBM
ICPC

出現するアルファベット

ABCIMP

係数つきの方程式

100*A + 10*B + 10*C + 100*I + 2*M = 1000*I + 101*C + 10*P

左辺から右辺を引く。この等式が成り立つようなアルファベットの値の組を全探索すれば良い事になります。

100*A + 10*B - 91*C - 900*I + 2*M - 10*P = 0

早速実装。ソースはとてもすっきり。

#include <iostream>
#include <set>
#include <vector>

using namespace std;

bool not_zero[10];
int coefficient[10];

int Dfs(int depth, int alphabet_size, int used, int lr) {
  if (depth == alphabet_size) {
    return lr == 0;
  } else {
    int ret = 0;
    for (int i = not_zero[depth] ? 1 : 0; i < 10; ++i) {
      if (used & (1 << i)) continue;
      ret += Dfs(depth + 1, alphabet_size, (used | (1 << i)), lr + coefficient[depth] * i);
    }
    return ret;
  }
}

int main() {
  int N;
  while (cin >> N, N) {
    vector<string> v(N);
    set<char> alph;
    for (int i = 0; i < N; ++i) {
      cin >> v[i];
      for (int j = 0; j < v[i].size(); ++j)
        alph.insert(v[i][j]);
    }

    string alphabet(alph.begin(), alph.end());

    vector<vector<int> > words(N);
    for (int i = 0; i < v.size(); ++i)
      for (int j = 0; j < v[i].size(); ++j)
        words[i].push_back(alphabet.find(v[i][j]));

    fill(coefficient, coefficient + 10, 0);
    for (int i = 0; i < N - 1; ++i)
      for (int j = words[i].size() - 1, ten = 1; j >= 0; --j, ten *= 10)
        coefficient[words[i][j]] += ten;
    for (int i = words.back().size() - 1, ten = 1; i >= 0; --i, ten *= 10)
      coefficient[words.back()[i]] -= ten;

    fill(not_zero, not_zero + 10, false);
    for (int i = 0; i < v.size(); ++i)
      if (v[i].size() > 1)
        not_zero[alphabet.find(v[i][0])] = true;

    cout << Dfs(0, alphabet.size(), 0, 0) << endl;
  }
}

55行になりました。Imos Judgeもこれで通ります。

さて、上のプログラムは制限時間15secのところ11.564secで通ったのですが、Judgeの実行履歴を見ると恐ろしい事に3sec程度で通している人たちがいます（しかも一発目で）。どうなってんだこれ、って事でみんなでしばらく考え、定数倍の最適化を施して何度か送信するも10secを切る事すら叶いませんでした。

その場は解散になったのですが、家に帰る途中に考えてみると…
・求めたいのは、アルファベットが変数となっている式の値が0になる組み合わせの数
・アルファベットを前半と後半に分けて、それぞれ組み合わせを全探索した後に和が0になる組を数えれば良さそう
・一つの組が持つのは(値, 使用した数字(ビットで持つ))で…
・使用した数字が重複せず、さらに値の和が0になるような組み合わせを数える
・ソートして尺取りでやればO(nlogn + n) = O(nlogn)で終わる
・（というかSRMで見かけたような…）

というわけで、実装は綺麗じゃないですが、O(nlogn) ( ただし、n = (N / 2)! )のプログラムになりました。以下ソース。

#include <iostream>
#include <set>
#include <map>
#include <vector>

using namespace std;

bool not_zero[10];
int coefficient[10];

struct S {
  int value, used;
  S(int value, int used) : value(value), used(used) {}
  bool operator<(const S& s) const {
    return value != s.value ? value < s.value : used < s.used;
  }
};

void Dfs(int depth, int alphabet_size, int used, int lr, map<S, int>& m) {
  if (depth == alphabet_size) {
    ++m[S(lr, used)];
  } else {
    for (int i = not_zero[depth] ? 1 : 0; i < 10; ++i) {
      if (used & (1 << i)) continue;
      Dfs(depth + 1, alphabet_size, used | (1 << i), lr + coefficient[depth] * i, m);
    }
  }
}

int main() {
  int N;
  while (cin >> N, N) {
    vector<string> v(N);
    set<char> alph;
    for (int i = 0; i < N; ++i) {
      cin >> v[i];
      for (int j = 0; j < v[i].size(); ++j)
        alph.insert(v[i][j]);
    }

    string alphabet(alph.begin(), alph.end());

    vector<vector<int> > words(N);
    for (int i = 0; i < v.size(); ++i)
      for (int j = 0; j < v[i].size(); ++j)
        words[i].push_back(alphabet.find(v[i][j]));

    fill(coefficient, coefficient + 10, 0);
    for (int i = 0; i < N - 1; ++i)
      for (int j = words[i].size() - 1, ten = 1; j >= 0; --j, ten *= 10)
        coefficient[words[i][j]] += ten;
    for (int i = words.back().size() - 1, ten = 1; i >= 0; --i, ten *= 10)
      coefficient[words.back()[i]] -= ten;

    fill(not_zero, not_zero + 10, false);
    for (int i = 0; i < v.size(); ++i)
      if (v[i].size() > 1)
        not_zero[alphabet.find(v[i][0])] = true;

    map<S, int> head;
    Dfs(0, alphabet.size() / 2, 0, 0, head);
    map<S, int> tail;
    Dfs(alphabet.size() / 2, alphabet.size(), 0, 0, tail);

    vector<pair<S, int> > h(head.begin(), head.end());
    vector<pair<S, int> > t(tail.rbegin(), tail.rend());

    int answer = 0;
    for (int i = 0, j = 0; i < h.size() && j < t.size();) {
      int hv = h[i].first.value, tv = t[j].first.value, value = hv + tv;
      if (value < 0) {
        ++i;
      } else if (value > 0) {
        ++j;
      } else {
        for (int k = i; k < h.size() && h[k].first.value == hv; ++k) {
          for (int l = j; l < t.size() && t[l].first.value == tv; ++l) {
            if ((h[k].first.used & t[l].first.used) == 0) {
              answer += h[k].second * t[l].second;
            }
          }
        }
        while (i < h.size() && h[i].first.value == hv) ++i;
        while (j < t.size() && t[j].first.value == tv) ++j;
      }
    }
    cout << answer << endl;
  }
}

Imos Judgeに送信してみると0.897sec。なかなか良さげです。