😊

三井住友信託銀行プログラミングコンテスト2019 | D - Lucky PIN

2 min read

問題

https://atcoder.jp/contests/sumitrust2019/tasks/sumitb2019_d

考えたこと

そのまま貪欲にやるとO(N^3)の時間がかかりTLEとなる。効率的な解法を考える必要がある。

入力されたSの中の数字がどの位置にあるか大事である。以下のようなTrie木のような木を作れば順序を保ちながら可能な暗証番号を記録できる。

暗証番号は3桁なので木の深さは3までと制限できる。よってノードの数は最大でも10^3程度であり間に合わせる事ができる。

コード

実装時のTips

  • 木は構造体をコピーさせないためにポインタで作る
#include <bits/stdc++.h>

#include <atcoder/all>

using namespace std;
using namespace atcoder;
using ll = long long;
using ld = long double;
using uint = unsigned int;
using ull = unsigned long long;
const int MOD = 1e9 + 7;

struct Node {
  map<char, Node *> children;
  int depth = 0;
};

void add(Node *node, char c) {
  if (node->depth == 3) {
    return;
  }
  for (auto m : node->children) {
    add(m.second, c);
  }
  if (node->children.find(c) == node->children.end()) {
    Node *child = new Node();
    child->depth = node->depth + 1;
    node->children[c] = child;
  }
}

unordered_set<string> s;

void search(Node *node, string cur) {
  if (node->depth == 3) {
    s.insert(cur);
    return;
  }
  for (auto m : node->children) {
    search(m.second, cur + m.first);
  }
}

// 木に順に入れていき最後に探索する
int main() {
  Node *root = new Node();
  ll n;
  cin >> n;
  for (int i = 0; i < n; i++) {
    char c;
    cin >> c;
    add(root, c);
  }
  search(root, "");
  cout << s.size() << endl;
}

別解

ACしたあとに解説をみてみたらいろんな解法があったので紹介する。

暗証番号を全探索

暗証番号は000から999までの1000個なので暗証番号をVとし、以下を1000回とけばいい。

  • Sから、N-3桁を消して暗証番号Vを作ることができるか

コードは以下となる。

#include <bits/stdc++.h>

#include <atcoder/all>
#include <boost/format.hpp>

using namespace std;
using namespace atcoder;
using ll = long long;
using ld = long double;
using uint = unsigned int;
using ull = unsigned long long;
const int MOD = 1e9 + 7;

int main() {
  ll n;
  string s;
  cin >> n >> s;

  ll ans = 0;
  for (int i = 0; i < 1000; i++) {
    string t = (boost::format("%03d") % i).str();
    ll si = 0, ti = 0;
    while (si < n && ti < 3) {
      if (s[si] == t[ti]) {
        si++;
        ti++;
      } else {
        si++;
      }
    }
    if (ti == 3) {
      ans++;
    }
  }
  cout << ans << endl;
}

参考

https://atcoder.jp/contests/sumitrust2019/tasks/sumitb2019_d/editorial
https://www.youtube.com/watch?v=vVOwtLuh60U&feature=youtu.be&t=3151&ab_channel=AtCoderLive