[HackerRank] プログラミングコンテスト復習02 [Week of Code 24:Happy Ladybugs]

Contest 2問目。

02: Happy Ladybugs

概要

ボードゲーム: Happy Ladybugs
てんとう虫は、隣が同じ色のてんとう虫であれば嬉しい。
なので、空いている場所へ移動して同じ色同士隣あおうとする。

RBY_YBR
てんとう虫をアルファベットで表す。同じアルファベットは同じ色のてんとう虫を意味している。_は空白地帯を表していて、てんとう虫はこの空白地帯で移動できる。
例えば
RBYBY_R
と、てんとう虫Bが移動する。この例は最終的に
BBRRYY_
と移動して全てのてんとう虫が隣り合い幸せな状態となる。

入力されたてんとう虫が全て幸せになれるならYESを、そうでなければNOを出力する。

方針

同じアルファベットが2文字以上存在
且つ
空いた場所が1箇所以上存在
すれば、てんとう虫は移動できて、同じ色同士集まることができる。

正規表現を利用して、アルファベットが2文字以上存在するか確認する。 てんとう虫を表す文字列bugsを並び替える。
set関数でユニークな文字列を得る。 -> b_set
並び替えた文字列に対して、各アルファベットが2文字以上存在するか正規表現でチェックし、set関数に収める。 -> r_set
空間をset関数に収める。 -> u_set

b_setから差分を取って、何も残らなかったら、てんとう虫は全て幸せである。
なにか残ったら幸せではない。

例: B_RRBR

u_set = set(['_'])
b_set = set(['R', 'B', '_'])
r_set = set(['R', 'B'])

b_set - r_set - u_set
  -> set([])

てんとう虫は隣り合える。

例: X_Y__X

u_set = set(['_'])
b_set = set(['Y', 'X', '_'])
r_set = set(['X', '_'])

b_set - r_set - u_set
  -> set(['Y'])

てんとう虫は隣り合えない。

コード

#!/bin/python
import re
def check(bugs):
    if '_' in bugs:
        bugs = ''.join(sorted(bugs))
    u_set = set('_')
    b_set = set(bugs)
    r_set = set(re.findall(r'(.)\1', bugs))
    if not bool(b_set - r_set - u_set):
        print 'YES'
    else:
        print 'NO'

for i in range(int(raw_input())):
    m = int(raw_input())
    check(raw_input())

結果

Result: Accept
Score: 20pt

考察

単純に並び替えれば同じアルファベットが並ぶだろうから、各アルファベットが2文字以上あればいいだろうと考えた。
for文で数え上げることも考えたが、直近で正規表現を使う項目を学んだので、利用した。
もう少し、スマートなやり方がある気がするが、全テストケースで通ったので改善はのちのち。