Longest Word Chain

This is a TwoSigma OA question.

[Problem Statement] You are given a library with n words (w[0], w[1], ...,w[n-1]). You choose a word from it, and in each step, remove one letter from this word only if doing so yields another word in the library. What is the longest possible chain of these removal steps?

Constraints: 1 <= n <= 50000 1 <= the length of each string in w <= 50 Each string is composed of lowercase ascii letters only The function should take in an array of strings w as its argument and should return a single integer representing the length of the longest chain of character removals possible.

Example input/output a b ba bca bda bdca

Calling the function on the input above should output 4. The longest possible chain is "bdca" -> "bca" -> "ba" -> "a".

#include <iostream>
#include <unordered_map>
#include <string>
#include <vector>
#include <algorithm>

using namespace std;

int helper(const string& s, unordered_map<string,int>& word_map){
    int longest = 1;
    string short_str;
    for(size_t i = 0; i < s.size(); ++i){
        short_str = s.substr(0, i) + s.substr(i + 1);
        if(word_map.count(short_str)){
            longest = max(longest, word_map[short_str] + 1);
        }
    }
    return longest;
}

int longest_chain(vector<string>& words){
    if(words.empty()) return 0;

    int n = words.size();

    sort(words.begin(), words.end(), [](const string& a, const string& b) {
      return a.length() < b.length();
    });

    unordered_map<string,int> word_map;
    word_map[words[0]] = 1;

    int longest = 1;
    for(const string& s : words){
        int curLen = helper(s, word_map);
        longest = max(longest, curLen);
        word_map[s] = curLen;
    }

    return longest;
}

int main() {
  cout<<"Hello" << endl;
  vector<string> input1({"a", "b", "ba", "bca", "bda", "bdca"});
  int res1 = longest_chain(input1);
  cout << res1 << endl;
  vector<string> input2({"a", "zxb", "ba", "bca", "bda", "bdca", "zxbe", "azxbe", "azxpbe"});
  cout << longest_chain(input2) << endl;
  return 0;
}

646. Maximum Length of Pair Chain