dnaSequence.cpp

#include <iostream>
#include <string>
#include <vector>
#include <numeric>
#include <algorithm>
#include <map>
#include <limits>
#include <cctype>

using namespace std;
using vi = vector<int>;
using ll = long long;

// Suffix Array and LCP array structure
struct SuffixArray {
    string s;
    vi sa, lcp;

    SuffixArray(const string &str = "") {
        if (!str.empty()) build(str);
    }

    void build(const string &str) {
        s = str + "\0"; // Append null character
        int n = s.size();
        sa.resize(n);
        vi rank(n);
        iota(sa.begin(), sa.end(), 0);

        for (int i = 0; i < n; i++) {
            rank[i] = s[i];
        }

        for (int k = 1; k < n; k *= 2) {
            auto cmp = [&](int i, int j) {
                if (rank[i] != rank[j]) return rank[i] < rank[j];
                int ri = i + k < n ? rank[i + k] : -1;
                int rj = j + k < n ? rank[j + k] : -1;
                return ri < rj;
            };
            sort(sa.begin(), sa.end(), cmp);

            vi new_rank(n);
            new_rank[sa[0]] = 0;
            for (int i = 1; i < n; i++) {
                new_rank[sa[i]] = new_rank[sa[i - 1]] + cmp(sa[i - 1], sa[i]);
            }
            rank = new_rank;
            if (rank[sa[n - 1]] == n - 1) break;
        }

        // LCP array (Kasai's algorithm)
        lcp.resize(n, 0);
        vi inv_sa(n);
        for (int i = 0; i < n; i++) inv_sa[sa[i]] = i;

        int k = 0;
        for (int i = 0; i < n - 1; i++) {
            if (inv_sa[i] == 0) {
                k = 0;
                continue;
            }
            int j = sa[inv_sa[i] - 1];
            while (i + k < n && j + k < n && s[i + k] == s[j + k]) k++;
            lcp[inv_sa[i]] = k;
            if (k > 0) k--;
        }
    }
};

// Forward declaration
struct CompressedSuffixTree;

// Node for the Compressed Suffix Tree
struct Node {
    int id;
    int parent;
    map<char, int> children;
    int start, end;
    int depth;
    CompressedSuffixTree* tree;

    Node(CompressedSuffixTree* t, int p, int start, int end, int depth)
        : id(-1), parent(p), start(start), end(end), depth(depth), tree(t) {}

    int len() const { return end - start; }
};

// Compressed Suffix Tree
struct CompressedSuffixTree {
    string S;
    SuffixArray sa_obj;
    vector<Node> nodes;
    int next_node_id = 0;

    CompressedSuffixTree(const string& str = "") {
        if (!str.empty()) {
            build(str);
        }
    }

    void build(const string& str) {
        S = str;
        sa_obj.build(S);
        nodes.clear();
        next_node_id = 0;

        // Create root node
        int root = new_node(-1, 0, 0, 0);

        // Build compressed suffix tree from suffix array and LCP (stack-based)
        // We skip the sentinel suffix (index == original length)
        int N = static_cast<int>(sa_obj.sa.size());
        int original_n = static_cast<int>(S.size());

        // Stack of node ids representing the current path; start with root
        vector<int> st;
        st.push_back(root);

        for (int i = 0; i < N; ++i) {
            int cur_sa = sa_obj.sa[i];
            // Skip sentinel suffix that points to the appended '\0'
            if (cur_sa == original_n) continue;

            int cur_lcp = (i == 0) ? 0 : sa_obj.lcp[i];

            // Pop until top.depth <= cur_lcp
            while (!st.empty() && nodes[st.back()].depth > cur_lcp) st.pop_back();
            int top = st.back();

            if (nodes[top].depth == cur_lcp) {
                // Create leaf as child of top
                int edge_start = cur_sa + nodes[top].depth;
                int edge_end = original_n; // exclusive
                int leaf = new_node(top, edge_start, edge_end, original_n - cur_sa);
                // Add child keyed by its first character
                char key = sa_obj.s[edge_start];
                nodes[top].children[key] = leaf;
                st.push_back(leaf);
            } else {
                // nodes[top].depth < cur_lcp -> need to create an internal node between top and its last child
                // The edge to split starts at SA[i-1] + nodes[top].depth
                int prev_sa = sa_obj.sa[i - 1];
                int split_start = prev_sa + nodes[top].depth;
                int split_end = prev_sa + cur_lcp;

                // Create new internal node
                int internal = new_node(top, split_start, split_end, cur_lcp);

                // The child that will be moved under the internal node is the node that was most recently added
                int child = st.back();
                // Record the child's current leading character (before we modify its start)
                char child_key = sa_obj.s[nodes[child].start];

                // Replace parent's child pointer to point to internal
                nodes[top].children[child_key] = internal;

                // Set internal as parent of child, and adjust child's edge start
                nodes[child].parent = internal;
                nodes[child].start = split_end;
                // Recompute child's depth (string-depth)
                nodes[child].depth = nodes[internal].depth + nodes[child].len();

                // Add the moved child under internal
                nodes[internal].children[sa_obj.s[nodes[child].start]] = child;

                // Finally add a new leaf for the current suffix under internal
                int edge_start = cur_sa + nodes[internal].depth;
                int edge_end = original_n;
                int leaf = new_node(internal, edge_start, edge_end, original_n - cur_sa);
                nodes[internal].children[sa_obj.s[edge_start]] = leaf;

                // Push internal and leaf onto stack
                st.push_back(internal);
                st.push_back(leaf);
            }
        }
    }

    int new_node(int parent, int start, int end, int depth) {
        nodes.emplace_back(this, parent, start, end, depth);
        nodes.back().id = next_node_id++;
        return nodes.back().id;
    }
    
    // Functionalities
    string longestRepeatedSubstring() {
        if (S.empty()) return "";
        
        int max_lcp = 0;
        int max_lcp_index = 0;
        for (size_t i = 1; i < sa_obj.lcp.size(); ++i) {
            if (sa_obj.lcp[i] > max_lcp) {
                max_lcp = sa_obj.lcp[i];
                max_lcp_index = sa_obj.sa[i];
            }
        }
        if (max_lcp > 0 && max_lcp_index + max_lcp <= static_cast<int>(S.length())) {
            return S.substr(max_lcp_index, max_lcp);
        }
        return "";
    }

    ll countDistinctSubstrings() {
        ll n = static_cast<ll>(S.length());
        ll total_substrings = n * (n + 1) / 2;
        ll redundant_substrings = 0;
        for (int l : sa_obj.lcp) {
            redundant_substrings += static_cast<ll>(l);
        }
        return total_substrings - redundant_substrings;
    }

    bool hasSubstring(const string& pattern) {
        if (pattern.empty()) return false;
        
        int m = static_cast<int>(pattern.length());
        int low = 0, high = static_cast<int>(sa_obj.sa.size()) - 1;
        
        while (low <= high) {
            int mid = low + (high - low) / 2;
            // Ensure we don't go out of bounds
            if (sa_obj.sa[mid] + m > static_cast<int>(S.length())) {
                high = mid - 1;
                continue;
            }
            
            int cmp = S.compare(sa_obj.sa[mid], m, pattern);
            if (cmp == 0) return true;
            if (cmp < 0) low = mid + 1;
            else high = mid - 1;
        }
        return false;
    }

    int countOccurrences(const string& pattern) {
        int m = pattern.length();
        int n = sa_obj.sa.size();

        auto cmp = [&](int suffixIndex) {
            int len = min(m, n - suffixIndex); // compare only available characters
            return S.substr(suffixIndex, len) < pattern;
        };

        // Find the first suffix >= pattern
        int low = 0, high = n;
        while (low < high) {
            int mid = low + (high - low) / 2;
            if (cmp(sa_obj.sa[mid])) low = mid + 1;
            else high = mid;
        }
        int first = low;

        // Check if pattern exists at all
        if (first == n || S.compare(sa_obj.sa[first], m, pattern) != 0) return 0;

        // Find the last suffix > pattern
        low = first, high = n;
        auto cmpUpper = [&](int suffixIndex) {
            int len = min(m, n - suffixIndex);
            return S.substr(suffixIndex, len) <= pattern;
        };
        while (low < high) {
            int mid = low + (high - low) / 2;
            if (cmpUpper(sa_obj.sa[mid])) low = mid + 1;
            else high = mid;
        }
        int last = low - 1;

        return last - first + 1;
    }
    
    string longestPalindromicSubstring() {
        int n = static_cast<int>(S.length());
        if (n == 0) return "";
        int start = 0, max_len = 1;

        for (int i = 0; i < n; ++i) {
            // Odd
            int l = i, r = i;
            while (l >= 0 && r < n && S[l] == S[r]) {
                if (r - l + 1 > max_len) {
                    max_len = r - l + 1;
                    start = l;
                }
                l--; r++;
            }
            // Even
            l = i; r = i + 1;
            while (l >= 0 && r < n && S[l] == S[r]) {
                if (r - l + 1 > max_len) {
                    max_len = r - l + 1;
                    start = l;
                }
                l--; r++;
            }
        }
        return S.substr(start, max_len);
    }


};

string longestCommonSubstring(const string& s1, const string& s2) {
    string combined = s1 + "#" + s2;
    SuffixArray sa(combined);
    int max_lcp = 0;
    int max_lcp_index = 0;
    int s1_len = static_cast<int>(s1.length());
    int separator_pos = s1_len; // Position of '#'

    for (size_t i = 1; i < sa.sa.size(); ++i) {
        // Check if suffixes come from different strings
        if ((sa.sa[i] <= separator_pos && sa.sa[i - 1] > separator_pos) ||
            (sa.sa[i - 1] <= separator_pos && sa.sa[i] > separator_pos)) {
            if (sa.lcp[i] > max_lcp) {
                max_lcp = sa.lcp[i];
                max_lcp_index = sa.sa[i];
            }
        }
    }
    if (max_lcp > 0) {
        return combined.substr(max_lcp_index, max_lcp);
    }
    return "";
}


void printMenu() {
    cout << "\n======================================\n";
    cout << "    DNA Sequence Analysis Tool\n";
    cout << "======================================\n";
    cout << "1. Analyze a single DNA sequence\n";
    cout << "2. Find Longest Common Substring (LCS) of two sequences\n";
    cout << "0. Exit\n";
    cout << "--------------------------------------\n";
    cout << "Enter your choice: ";
    cout.flush(); 
}

void printSingleStringMenu() {
    cout << "\n--- Single Sequence Analysis ---\n";
    cout << "1. Find Longest Repeated Substring\n";
    cout << "2. Count Distinct Substrings\n";
    cout << "3. Search for a Pattern\n";
    cout << "4. Count Pattern Occurrences\n";
    cout << "5. Find Longest Palindromic Substring\n";
    cout << "0. Back to main menu\n";
    cout << "------------------------------\n";
    cout << "Enter your choice: ";
    cout.flush(); 
}

bool isValidDNASequence(const string& seq) {
    if (seq.empty()) return false;
    for (char c : seq) {
        if (c != 'A' && c != 'T' && c != 'G' && c != 'C' && 
            c != 'a' && c != 't' && c != 'g' && c != 'c') {
            return false;
        }
    }
    return true;
}

void handleSingleStringAnalysis() {
    string s;
    cout << "Enter a DNA sequence: ";
    cout.flush();
    cin >> s;
    if (s.empty()) {
        cout << "Error: DNA sequence cannot be empty." << endl;
        return;
    }
    
    // Convert to uppercase for consistency
    transform(s.begin(), s.end(), s.begin(), ::toupper);
    
    if (!isValidDNASequence(s)) {
        cout << "Error: Invalid DNA sequence. Only A, T, G, C characters are allowed." << endl;
        return;
    }

    CompressedSuffixTree cst(s);

    int choice;
    do {
        printSingleStringMenu();
        
        // Clear any leftover input
        if (cin.peek() == '\n') {
            cin.ignore();
        }
        
        cin >> choice;
        
        if (cin.fail()) {
            cin.clear();
            cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
            cout << "Invalid input. Please enter a number." << endl;
            choice = -1; // Invalid choice, retry
            continue;
        }
        
        // Clear any remaining characters in the input buffer
        cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
        switch (choice) {
            case 1: {
                string lrs = cst.longestRepeatedSubstring();
                if (!lrs.empty()) {
                    cout << "Longest Repeated Substring: \"" << lrs << "\"" << endl;
                } else {
                    cout << "No repeated substring found." << endl;
                }
                break;
            }
            case 2:
                cout << "Number of distinct substrings: " << cst.countDistinctSubstrings() << endl;
                break;
            case 3: {
                string p;
                cout << "Enter pattern to find: ";
                cout.flush();
                cin >> p;
                if (p.empty()) {
                    cout << "Error: Pattern cannot be empty." << endl;
                    break;
                }
                if (cst.hasSubstring(p)) {
                    cout << "Pattern \"" << p << "\" found in the sequence." << endl;
                } else {
                    cout << "Pattern \"" << p << "\" not found in the sequence." << endl;
                }
                break;
            }
            case 4: {
                string p;
                cout << "Enter pattern to count: ";
                cout.flush();
                cin >> p;
                if (p.empty()) {
                    cout << "Error: Pattern cannot be empty." << endl;
                    break;
                }
                int count = cst.countOccurrences(p);
                cout << "Pattern \"" << p << "\" occurs " << count << " time(s) in the sequence." << endl;
                break;
            }
            case 5: {
                string lps = cst.longestPalindromicSubstring();
                if (!lps.empty()) {
                    cout << "Longest Palindromic Substring: \"" << lps << "\" (length: " << lps.length() << ")" << endl;
                } else {
                    cout << "No palindrome found." << endl;
                }
                break;
            }
            case 0:
                cout << "Returning to main menu..." << endl;
                break;
            default:
                cout << "Invalid choice. Please select 0-5." << endl;
                break;
        }
    } while (choice != 0);
}

int main() {
    // Don't disable sync for better interactive behavior
    // ios::sync_with_stdio(false);
    // cin.tie(nullptr);

    cout << "Welcome to the DNA Sequence Analysis Tool!" << endl;
    cout << "This tool uses Compressed Suffix Trees for efficient analysis." << endl;
    cout.flush(); // Ensure output is displayed immediately

    int choice;
    do {
        printMenu();
        
        // Clear any leftover input
        if (cin.peek() == '\n') {
            cin.ignore();
        }
        
        cin >> choice;
        
        if (cin.fail()) {
            cin.clear();
            cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
            cout << "Invalid input. Please enter a number." << endl;
            continue;
        }
        
        // Clear any remaining characters in the input buffer
        cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
        switch (choice) {
            case 1:
                handleSingleStringAnalysis();
                break;
            case 2: {
                string s1, s2;
                cout << "Enter first DNA sequence: ";
                cout.flush();
                cin >> s1;
                cout << "Enter second DNA sequence: ";
                cout.flush();
                cin >> s2;
                
                if (s1.empty() || s2.empty()) {
                    cout << "Error: Both DNA sequences must be non-empty." << endl;
                    continue;
                }
                
                // Convert to uppercase for consistency
                transform(s1.begin(), s1.end(), s1.begin(), ::toupper);
                transform(s2.begin(), s2.end(), s2.begin(), ::toupper);
                
                if (!isValidDNASequence(s1) || !isValidDNASequence(s2)) {
                    cout << "Error: Invalid DNA sequence(s). Only A, T, G, C characters are allowed." << endl;
                    continue;
                }
                
                string lcs = longestCommonSubstring(s1, s2);
                if (!lcs.empty()) {
                    cout << "Longest Common Substring: \"" << lcs << "\" (length: " << lcs.length() << ")" << endl;
                } else {
                    cout << "No common substring found between the sequences." << endl;
                }
                break;
            }
            case 0:
                cout << "\nThank you for using the DNA Sequence Analysis Tool!" << endl;
                cout << "Goodbye!" << endl;
                break;
            default:
                cout << "Invalid choice. Please enter 0, 1, or 2." << endl;
                break;
        }
    } while (choice != 0);

    return 0;
}