tmdb_core.cpp

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <vector>
#include <string>
#include <algorithm>
#include <unordered_set>

namespace py = pybind11;

int levenshtein_distance( const std::string& s1, const std::string& s2 ) {
    if( s1.size() < s2.size() ) {
        return levenshtein_distance( s2, s1 );
    }
    
    std::vector<int> prev( s2.size() + 1 );
    std::vector<int> curr( s2.size() + 1 );
    
    for( size_t j = 0; j <= s2.size(); ++j ) {
        prev[ j ] = j;
    }

    for( size_t i = 1; i <= s1.size(); ++i ) {
        curr[ 0 ] = i;
        for( size_t j = 1; j <= s2.size(); ++j ) {
            int cost = ( s1[ i - 1 ] == s2[ j - 1 ] ) ? 0 : 1;
            curr[ j ] = std::min( { prev[ j ] + 1, curr[ j - 1 ] + 1, prev[ j - 1 ] + cost } );
        }
        prev = curr; 
    }
    return prev[ s2.size() ];
}

std::vector<int> get_actor_overlap( const std::vector<int>& cast1, const std::vector<int>& cast2 ) {
    std::unordered_set<int> set1( cast1.begin(), cast1.end() );
    std::vector<int> overlap;
    overlap.reserve( std::min( cast1.size(), cast2.size() ) ); 

    for( int actor_id : cast2 ) {
        if( set1.find( actor_id ) != set1.end() ) {
            overlap.emplace_back( actor_id );
        }
    }
    return overlap;
}

PYBIND11_MODULE( tmdb_core, m ) {
    m.def( "levenshtein_distance", &levenshtein_distance );
    m.def( "get_actor_overlap", &get_actor_overlap );
}