AtCoder C - Neo-lexicographic OrderingがPythonで再現できない

前提・実現したいこと

AtCoderの過去問を解いて練習しているのですが、
以下の問題が解説を読んでも理解できずにいます。
https://atcoder.jp/contests/abc219/tasks/abc219_c

発生している問題・エラーメッセージ

WAになってしまう。

該当のソースコード

以下のソースを作成しました。

python
1import functools
2
3X: list = list(input())
4pos: dict = {}
5
6for i, x in enumerate(X):
7    pos[x] = i
8
9N: int = int(input())
10
11S: list = []
12for _ in range(N):
13    S.append(input())
14
15
16def original_sort(s1, s2):
17    length = min(len(s1), len(s2))
18    if len(s1) > len(s2):
19        s2, s1 = s1, s2
20
21    count = 1
22    for i in range(length):
23        if pos[s1[i]] > pos[s2[i]]:
24            return 1
25        elif pos[s1[i]] < pos[s2[i]]:
26            return -1
27        else:
28            if count == length:
29                return 0
30            count += 1
31
32
33sorted_S = sorted(S, key=functools.cmp_to_key(original_sort))
34
35[print(s) for s in sorted_S]
36

試したこと

公式解説がC++なので、Pythonで再現できているか不明です。
どの辺りがおかしいかご教示いただきたいです。
よろしくお願いいたします。

行動規範の内容に同意します

回答2件

どの辺りがおかしいか

質問に掲載されている Python スクリプトに、例えば以下のデータを入力すると、

data
1zyxwvutsrqponmlkjihgfedcba
22
3zabc
4z

結果は以下になります。

zabc
z

おかしいのは以下の if count == length: return 0 の部分です。

python
1def original_sort(s1, s2):
2    length = min(len(s1), len(s2))
3    if len(s1) > len(s2):
4        s2, s1 = s1, s2
5
6    count = 1
7    for i in range(length):
8        if pos[s1[i]] > pos[s2[i]]:
9            return 1
10        elif pos[s1[i]] < pos[s2[i]]:
11            return -1
12        else:
13            if count == length:
14                return 0
15            count += 1

追記

元のコードを書き直す場合は original_sort を以下の様にします。

python
1def original_sort(s1, s2):
2    length = min(len(s1), len(s2))
3    for i in range(length):
4        if pos[s1[i]] > pos[s2[i]]:
5            return 1
6        elif pos[s1[i]] < pos[s2[i]]:
7            return -1
8    return 2*(len(s1) > len(s2)) - 1

投稿2021/12/13 12:29

編集2021/12/14 00:57

melian

総合スコア19865

utti_keima

2021/12/13 22:43

ご回答いただき、ありがとうございます！確かにおっしゃる通りですね、、考え直してみます。

行動規範の内容に同意します

ベストアンサー

どこがおかしいか？というよりも公式解説の通りに実装すればよいだけかと思います。

Python
1def original_sort(s1, s2):
2    length = min(len(s1), len(s2))
3    for i in range(length):
4        if s1[i] != s2[i]:
5            return pos[s1[i]] - pos[s2[i]]
6    return len(s1) - len(s2)

ちなみにPythonだとtupleがソート可能なので以下のように短く書けます。

Python
1X = list(input())
2pos = {x:i for i, x in enumerate(X)}
3N = int(input())
4S = [input() for _ in range(N)]
5S = sorted([tuple([pos[c] for c in s]) for s in S])
6[print(''.join([X[c] for c in s])) for s in S]