Problem statement

https://binarysearch.com/problems/Word-Distance-Queries/

Solution

Equal to Leetcode 0244 Shortest Word Distance II.

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 WordDistanceQuerier:
    def __init__(self, wordsDict):
        self.places = defaultdict(list)
        for i, word in enumerate(wordsDict):
            self.places[word].append(i)

    def distance(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)