서로 다른 고유한 단어들로 이루어진 집합이 주어졌을 때, 이 단어들을 조합하여 만들 수 있는 모든 단어 사각형(Word Square)을 찾아야 합니다. 여기서 단어 사각형이란 k번째 행과 k번째 열이 정확히 동일한 문자열이 되는 단어 시퀀스를 의미하며, 조건은 0 ≤ k < max(numRows, numColumns)입니다.
예를 들어 ["ball", "area", "lead", "lady"]라는 단어 시퀀스는 아래와 같이 배치했을 때 가로 방향과 세로 방향 어느 쪽으로 읽어도 같은 문자열이 되므로 유효한 단어 사각형을 구성합니다.
| b | a | l | l |
| a | r | e | a |
| l | e | a | d |
| l | a | d | y |
따라서 입력이 ["area", "lead", "wall", "lady", "ball"]로 주어진다면, 출력은 [["wall", "area", "lead", "lady"], ["ball", "area", "lead", "lady"]]가 됩니다.
문제 해결 접근 방식
이 문제는 트라이(Trie) 자료구조와 백트래킹(Backtracking)을 결합하면 효율적으로 해결할 수 있습니다. 이미 배치된 단어들로부터 다음 행의 첫 글자들이 만드는 접두사(prefix)를 구한 뒤, 트라이에서 해당 접두사로 시작하는 단어만 빠르게 조회하는 것이 핵심입니다. 전체 과정은 다음과 같습니다.
1. 노드 구조 정의
- 단어의 끝을 표시하는 isEnd 변수와 자식 노드들을 저장하는 child 맵을 가지는 노드 구조체를 정의합니다.
- 결과를 담을 2차원 배열 ret을 선언합니다.
2. insertNode() — 트라이에 단어 삽입
- head와 문자열 s를 매개변수로 받습니다.
- node := head로 초기화한 뒤, i := 0부터 s의 길이 미만까지 반복합니다.
- x := s[i]
- node의 child에 x에 해당하는 노드가 없다면 새 노드를 생성하여 child[x]에 할당합니다.
- node := node의 child[x]로 이동합니다.
- 반복이 끝나면 마지막 노드의 isEnd를 true로 설정합니다.
3. getAllWords() — 접두사와 일치하는 단어 수집
- idx, prefix, node, temp 배열을 매개변수로 받습니다.
- node가 비어 있으면 즉시 반환합니다.
- node의 isEnd가 true이면 현재까지 만든 문자열 curr을 temp 끝에 삽입하고 반환합니다.
- idx가 prefix의 길이 이상이면, node의 모든 자식 노드에 대해 getAllWords(idx, prefix, it.second, temp, curr + it.first)를 재귀 호출합니다.
- 그렇지 않은 경우 x := prefix[idx]로 설정하고, node의 child[x]가 존재하지 않으면 반환합니다. 존재한다면 getAllWords(idx + 1, prefix, child[x], temp, curr + x)를 호출합니다.
4. solve() — 백트래킹으로 사각형 확장
- temp 배열, idx, reqSize, head를 매개변수로 받습니다.
- idx == reqSize이면 하나의 단어 사각형이 완성된 것이므로 temp를 ret에 추가하고 반환합니다.
- prefix := 빈 문자열로 초기화한 뒤, i := 0부터 temp의 크기 미만까지 반복하며 prefix += temp[i][idx]로 다음 행이 가져야 할 접두사를 만듭니다.
- possible 배열을 정의하고 curr = head로 설정한 후 getAllWords(0, prefix, curr, possible)를 호출하여 후보 단어를 수집합니다.
- possible의 각 단어 s에 대해 temp에 s를 추가하고 solve(temp, idx + 1, reqSize, head)를 재귀 호출한 뒤, 마지막 요소를 제거하여 백트래킹합니다.
5. 메인 로직
- head = new node로 트라이의 루트 노드를 생성합니다.
- 모든 단어에 대해 insertNode(head, words[i])를 호출하여 트라이를 구축합니다.
- 임시 배열 temp를 정의하고, 각 단어 s := words[i]를 temp에 넣은 뒤 solve(temp, 1, words[0].size(), head)를 호출하고 다시 제거합니다. 첫 번째 단어를 무엇으로 시작하느냐에 따라 모든 경우를 탐색하기 위함입니다.
- 최종적으로 ret을 반환합니다.
예제 코드
아래 C++ 구현을 통해 더 자세히 이해해 보겠습니다 −
#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<vector<auto>> v){
cout << "[";
for(int i = 0; i<v.size(); i++){
cout << "[";
for(int j = 0; j <v[i].size(); j++){
cout << v[i][j] << ", ";
}
cout << "],";
}
cout << "]"<<endl;
}
struct Node {
bool isEnd;
map<char, Node *> child;
};
class Solution {
public:
vector<vector<string>> ret;
void insertNode(Node *head, string &s) {
Node *node = head;
for (int i = 0; i < s.size(); i++) {
char x = s[i];
if (!node->child[x]) {
node->child[x] = new Node();
}
node = node->child[x];
}
node->isEnd = true;
}
void getAllWords(int idx, string prefix, Node *node, vector<string>&temp,
string curr = "") {
if (!node)
return;
if (node->isEnd) {
temp.push_back(curr);
return;
}
if (idx >= prefix.size()) {
for (auto &it : node->child) {
getAllWords(idx, prefix, it.second, temp, curr + it.first);
}
}
else {
char x = prefix[idx];
if (!node->child[x])
return;
getAllWords(idx + 1, prefix, node->child[x], temp, curr + x);
}
}
void solve(vector<string> &temp, int idx, int reqSize, Node *head){
if (idx == reqSize) {
ret.push_back(temp);
return;
}
string prefix = "";
for (int i = 0; i < temp.size(); i++) {
prefix += temp[i][idx];
}
vector<string> possible;
Node *curr = head;
getAllWords(0, prefix, curr, possible);
for (int i = 0; i < possible.size(); i++) {
string s = possible[i];
temp.push_back(s);
solve(temp, idx + 1, reqSize, head);
temp.pop_back();
}
}
vector<vector<string>> wordSquares(vector<string> &words) {
ret.clear();
Node *head = new Node();
for (int i = 0; i < words.size(); i++) {
insertNode(head, words[i]);
}
vector<string> temp;
for (int i = 0; i < words.size(); i++) {
string s = words[i];
temp.push_back(s);
solve(temp, 1, (int)words[0].size(), head);
temp.pop_back();
}
return ret;
}
};
main() {
Solution ob;
vector<string> v = {"area", "lead", "wall", "lady", "ball"};
print_vector(ob.wordSquares(v));
}입력
{"area", "lead", "wall", "lady", "ball"}출력
[[wall, area, lead, lady],[ball, area, lead, lady]]