Problem statement

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

Solution

If words are not equal, use problem 0243. If they are equal, it is again two pointers, but when see new occurrence of the word, we update the smallest index

Complexity

Time complexity is $O(n)$ and space complexity is $O(1)$.

Code

class Solution:
    def shortestWordDistance(self, wordsDict, word1, word2):
        p1, p2 = -float("inf"), float("inf")
        ans = float("inf")
        
        if word1 != word2:
            for i, word in enumerate(wordsDict):
                if word == word1:
                    p1 = i
                    ans = min(ans, abs(p2 - p1))
                if word == word2:
                    p2 = i
                    ans = min(ans, abs(p2 - p1))
        else:
            for i, word in enumerate(wordsDict):
                if word == word1:
                    p1, p2 = p2, i
                    ans = min(ans, abs(p2 - p1))
                    
        return ans