Problem statement

https://leetcode.com/problems/shortest-word-distance-ii/

Solution

Create hash table, which for each word have all indexes for this word in original list. Then when we given two words, we need to look into 2 corresponding values for our hash table, which are lists, and compare lists in $O(a+b)$, where $a$ and $b$ are the length of these lists — idea similar to merge sort.

In py code I use the trick that python uses timsort, which merge two sorted lists in linear time.

Complexity

It is $O(n)$ to consturct data structure and $O(a+b)$, where $a$ and $b$ occurunces of words for shortest function. Note, that it can also be done in $O(a\log b)$ or $O(b\log a)$ complexity, if we use binary search.

Code

class WordDistance:
    def __init__(self, wordsDict):
        self.places = defaultdict(list)
        for i, word in enumerate(wordsDict):
            self.places[word].append(i)

    def shortest(self, word1, word2):
        p1, p2 = self.places[word1], self.places[word2]
        t = sorted([(i1, 1) for i1 in p1] + [(i2, 2) for i2 in p2])
        return min(j[0]-i[0] for i,j in zip(t, t[1:]) if j[1] + i[1] == 3)