12월 13일 - 문자열 집합 판별[BOJ/9250]

Yullgiii·2024년 12월 13일

개발자로의 취업을 준비하며 TIL

목록 보기

333/363

문제 설명

주어진 문자열 집합 𝑆의 각 문자열과 쿼리 문자열의 모든 부분 문자열이 일치하는지 판별하는 문제이다.
Aho-Corasick 알고리즘을 사용해 효율적으로 𝑄개의 문자열 쿼리를 처리해야 한다.
문제는 트라이(Trie)와 실패 함수(Failure Function)를 활용해 빠르게 부분 문자열을 탐색하고, 효율적으로 결과를 판별하는 데 중점을 둔다.

코드

import java.util.*;

class Node {
    boolean valid;
    Node[] child;
    Node fail;

    Node() {
        valid = false;
        child = new Node[26];
        fail = null;
    }

    void insert(String s, int idx) {
        if (idx == s.length()) {
            valid = true;
            return;
        }
        int x = s.charAt(idx) - 'a';
        if (child[x] == null) {
            child[x] = new Node();
        }
        child[x].insert(s, idx + 1);
    }
}

public class Main {
    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);
        int n = scanner.nextInt();
        scanner.nextLine(); // Consume newline

        Node root = new Node();

        for (int i = 0; i < n; i++) {
            String s = scanner.nextLine();
            root.insert(s, 0);
        }

        // Build Aho-Corasick automaton
        Queue<Node> queue = new LinkedList<>();
        queue.add(root);
        root.fail = root;

        while (!queue.isEmpty()) {
            Node current = queue.poll();

            for (int i = 0; i < 26; i++) {
                Node next = current.child[i];
                if (next == null) continue;

                if (current == root) {
                    next.fail = root;
                } else {
                    Node dest = current.fail;
                    while (dest != root && dest.child[i] == null) {
                        dest = dest.fail;
                    }
                    if (dest.child[i] != null) {
                        dest = dest.child[i];
                    }
                    next.fail = dest;
                }

                if (next.fail.valid) {
                    next.valid = true;
                }
                queue.add(next);
            }
        }

        int m = scanner.nextInt();
        scanner.nextLine(); // Consume newline

        for (int i = 0; i < m; i++) {
            String query = scanner.nextLine();

            Node current = root;
            boolean found = false;

            for (int j = 0; j < query.length(); j++) {
                int x = query.charAt(j) - 'a';
                while (current != root && current.child[x] == null) {
                    current = current.fail;
                }
                if (current.child[x] != null) {
                    current = current.child[x];
                }
                if (current.valid) {
                    found = true;
                    break;
                }
            }

            System.out.println(found ? "YES" : "NO");
        }

        scanner.close();
    }
}

코드 설명

1. 노드 구조 정의

class Node {
    boolean valid; // 이 노드가 집합 S의 문자열 끝인지 여부
    Node[] child; // 각 알파벳에 대한 자식 노드
    Node fail; // 실패 함수 포인터

    Node() {
        valid = false;
        child = new Node[26];
        fail = null;
    }

    void insert(String s, int idx) {
        if (idx == s.length()) {
            valid = true; // 문자열의 끝을 표시
            return;
        }
        int x = s.charAt(idx) - 'a';
        if (child[x] == null) {
            child[x] = new Node();
        }
        child[x].insert(s, idx + 1); // 다음 문자를 삽입
    }
}

Node 클래스는 각 트라이 노드의 상태를 나타낸다.
insert 메서드는 문자열을 트라이 구조에 삽입한다.

2. Aho-Corasick 자동화 구성

// Aho-Corasick 트라이 및 실패 함수 구축
Queue<Node> queue = new LinkedList<>();
queue.add(root);
root.fail = root;

while (!queue.isEmpty()) {
    Node current = queue.poll();

    for (int i = 0; i < 26; i++) {
        Node next = current.child[i];
        if (next == null) continue;

        if (current == root) {
            next.fail = root; // 루트의 자식은 항상 루트를 가리킴
        } else {
            Node dest = current.fail;
            while (dest != root && dest.child[i] == null) {
                dest = dest.fail;
            }
            if (dest.child[i] != null) {
                dest = dest.child[i];
            }
            next.fail = dest;
        }

        // 실패 노드가 유효한 경우, 이 노드도 유효함
        if (next.fail.valid) {
            next.valid = true;
        }
        queue.add(next);
    }
}

실패 함수는 트라이를 따라가며 다음 가능한 매칭 지점으로 이동하도록 구성한다.
BFS를 통해 실패 포인터를 계산하며, 유효한 노드인지도 업데이트한다.

3. 문자열 쿼리 처리

for (int i = 0; i < m; i++) {
    String query = scanner.nextLine();

    Node current = root;
    boolean found = false;

    for (int j = 0; j < query.length(); j++) {
        int x = query.charAt(j) - 'a';
        while (current != root && current.child[x] == null) {
            current = current.fail; // 실패 포인터 따라가기
        }
        if (current.child[x] != null) {
            current = current.child[x]; // 다음 상태로 이동
        }
        if (current.valid) { // 집합 S의 문자열과 매칭
            found = true;
            break;
        }
    }

    System.out.println(found ? "YES" : "NO");
}