Search Suggestions

/*!
 * \brief  Search suggestions.
 *
 * \author Alexey Medveschek
 * \copy   2015 (C) Go.Mail.Ru
 */

#include <iostream>
#include <fstream>
#include <sstream>
#include <map>
#include <vector>
#include <string>

typedef unsigned int id_t;           //!< индекс подсказки
typedef std::vector<id_t> idvec_t;   //!< списки (posting lists); будем пересекать при поиске

const id_t kInvalidId = (id_t)(~0U); //!< специальное обозначение "пустого" id

//-----------------------------------------------
//! Обратный индекс, он же трай, он же узел трая
struct InvertedIndex {
    typedef std::map<char, InvertedIndex> trans_t; //!< таблица переходов
    trans_t trans_;

    idvec_t ids_;

    //! Добавить слово в обратный индекс
    void add_word(const std::string& word, id_t id) {
        InvertedIndex* node = this;

        id_t last_id = kInvalidId;

        std::string::const_iterator wi = word.begin(), wend = word.end();
        for (; wi != wend; ++wi) {
            // на каждый символ в каждую ноду добавляем id подсказки
            node = &node->trans_[ *wi ];
            // добавляем только уникальные id
            last_id = (!node->ids_.empty() ? node->ids_[node->ids_.size()-1] : kInvalidId);
            if (last_id != id) {
                node->ids_.push_back( id );
            }
        }
    }

    //! Найти ноду по префиксу; NULL - если строка не найдена
    const InvertedIndex* find(const std::string& prefix) const {
        const InvertedIndex* node = this;

        std::string::const_iterator pi = prefix.begin(), pend = prefix.end();
        for (; pi != pend; ++pi) {
            trans_t::const_iterator ni = node->trans_.find( *pi );
            if (ni == node->trans_.end()) {
                return NULL;
            }

            node = &ni->second;
        }

        return node;
    }
};

InvertedIndex inverted_index;

//-----------------------------------------------
//! Прямой индекс
struct Sugg {
    std::string text_;
    unsigned int line_; //!< Вместо веса храним номер строки в файле
    Sugg(const std::string& text, unsigned line): text_(text), line_(line) {}
};

std::vector<Sugg> forward_index;

//-----------------------------------------------
// Парсит из входного потока подскакзи и добавляет их в индекс.
void build_index(std::istream& in) {
    std::string line, word;

    unsigned int line_num = 0;
    while (std::getline(in, line)) {
        // добавляем подсказку в прямой индекс
        const id_t id = forward_index.size();
        forward_index.push_back( Sugg(line, ++line_num) );

        // теперь в обратный
        std::stringstream s_line(line);
        while (std::getline(s_line, word, ' ')) {
            if (word.empty()) continue; // пропускаем двойные пробелы
            inverted_index.add_word(word, id);
        }
    }
}

//-----------------------------------------------
// "Ленивый" алгоритм пересечения сортированных по неубыванию списков;
// пересекает до получения TOP-K первых результатов.
// Возвращает true, если пересечение не пустое.
bool intersect_sorted_lists(const std::vector<const idvec_t*>& lists, size_t top, idvec_t & res) {
    res.clear();

    const size_t lnum = lists.size();
    if (! lnum)         // нет списков для пересечения
        return false;

    if (lnum == 1) {    // единственный список
        res = *lists[0];
        top = std::min(top, res.size());
        res.resize( top );
        return (! res.empty());
    }

    // подготавливаем пары iterator-end
    std::vector<std::pair<const id_t*, const id_t*> > lptrs(lnum);
    for (size_t i = 0; i < lnum; ++i) {
        const idvec_t & idvec = *lists[i];

        if (idvec.empty())      // один из списков пустой
            return false;

        lptrs[i].first  = &idvec[0];
        lptrs[i].second = &idvec[0] + idvec.size();
    }

    size_t li = 0, found = 0;
    id_t cur_id = *(lptrs[0].first);

    // идём, пока один из списков не закончился и пока не набрали нужное кол-во результатов
    while (res.size() < top) {
        // ищем текущий элемент в списке
        while (lptrs[li].first < lptrs[li].second && *lptrs[li].first < cur_id) {
            ++lptrs[li].first;
        }
        // один из списков закончился
        if (lptrs[li].first == lptrs[li].second) {
            break;
        }

        const id_t check_id = *(lptrs[li].first);

        if (check_id == cur_id) {
            ++found;
            if (found == lnum) {        // искомый элемент нашёлся во всех списках
                res.push_back( cur_id );
                // начинаем поиск нового элемента
                ++lptrs[li].first;
                if (lptrs[li].first < lptrs[li].second) {
                    cur_id = *lptrs[li].first;
                    found = 1;
                } else {
                    break;              // один из списков закончился
                }
            }
        } else {
            // начинаем поиск нового элемента
            cur_id = check_id;
            found = 1;
        }

        li = (li + 1) % lnum;           // на след. список
    }

    return (! res.empty());
}

//-----------------------------------------------
// Поиск подсказок по заданному запросу
// возвращает top первых самых весомых подсказок, вектор пар: подсказка-вес
bool search(const std::string& q, size_t top, std::vector<Sugg>& res) {
    res.clear();

    // разбиваем запрос на слова-префиксы, получаем списки для пересечения
    std::vector<const idvec_t*> lists;

    std::string word;
    std::stringstream stext(q);
    while (std::getline(stext, word, ' ')) {
        if (word.empty()) continue; // пропускаем двойные пробелы
        const InvertedIndex* node = inverted_index.find(word);
        if (! node)
            return false;
        lists.push_back( &node->ids_ );
    }

    if (lists.empty()) return false;

    // пересекаем списки
    idvec_t intersection;
    if (! intersect_sorted_lists(lists, top, intersection))
        return false;

    // извлекаем данные из прямого индекса
    for (size_t i = 0; i < intersection.size(); ++i) {
        const id_t id = intersection[i];
        res.push_back( forward_index[id] );
    }

    return true;
}

//-----------------------------------------------
int main(int argc, char* argv[]) {
    if (argc < 2) {
        std::cerr << "Usage: " << argv[0] << " <line_by_line_suggestions_file.txt>" << std::endl;
        return 1;
    }

    std::cerr << "Adding suggestions to index..." << std::endl;
    std::ifstream in(argv[1]);
    build_index(in);
    std::cerr << "Successfully added " << forward_index.size() << " suggestions\n";

    std::cerr << "Ready. Type any incomplete query:" << std::endl;

    std::string query;
    const size_t top = 10;
    std::vector<Sugg> res;

    while (std::getline(std::cin, query)) {
        if (! search(query, top, res)) {
            std::cout << "\t<nothing to suggest>\n" << std::endl;
            continue;
        }

        std::cout << query << std::endl;
        for (size_t i = 0; i < res.size(); ++i) {
            std::cout << "\t" << res[i].line_ << ". "
                      << res[i].text_ << std::endl;
        }
        std::cout << "\n";
    }

    return 0;
}