Code Formula 2014 予選B C - 仲良し文字列

問題原文

atcoder.jp

解法

3回のswapで影響範囲があるのは最大６文字まで。それより大きいと３回では足りない。
よって と が異なるindexを列挙しておいて、実際に３回のswapを全探索する。

の中に同じ文字が複数あれば、それらをswapすることで回数を稼げることに注目。
具体的にはそのような場合、 ちょうど３回のswap でなく ３回以下のswap と条件が変わる。

実装

こういうのは再帰で書くのが楽なのではないでしょうか。

from itertools import combinations
A = list(input())
B = list(input())

diff = []
dupulicate_char_index = -1
appeared = set()
for i, (a, b) in enumerate(zip(A, B)):
    if a != b:
        diff.append(i)
    if a in appeared:
        dupulicate_char_index = i
    appeared.add(a)
if dupulicate_char_index != -1:
    diff += [dupulicate_char_index] * (6 - len(diff))

# 達成不可能
if len(diff) > 6:
    print('NO')
    exit()


def check(a, change_cnt):
    if (dupulicate_char_index != -1 or change_cnt == 3) and (a == B):
        print('YES')
        exit()


def dfs(a, depth):
    if depth == 3:
        return
    for i, j in combinations(diff, 2):
        a[i], a[j] = a[j], a[i]
        check(a, depth + 1)
        dfs(a, depth + 1)
        a[i], a[j] = a[j], a[i]

dfs(A, 0)
print('NO')

感想

結構難しかった。。。