数组中两个字符串的最小距离

题目

给定一个字符串数组arr，再给定两个字符串s1, s2，返回在arr中，s1与s2的最小距离。如果s1, s2为None或不在arr中，返回-1。
例如：
arr = [‘1’, ‘3’, ‘3’, ‘3’, ‘2’, ‘3’, ‘1’]
s1 = ‘1’
s2 = ‘2’
返回2

arr = [‘CD’]
s1 = ‘CD’
s2 = ‘AB’
返回-1

思路

用pos1, pos2记录s1, s2出现的最后位置，如果遍历位置i找到s2且前面已经找到s1，则i-pos1为s1, s2的一个距离，记录找到的最小距离，每找到一个合适的位置，跟最小距离比较。

实现

def min_distance(arr, s1, s2):
    if arr is None or len(arr) == 0 or s1 is None or s2 is None:
        return -1
    if s1 == s2:
        return 0

    pos1, pos2 = -1, -1
    ans = sys.maxsize
    for i in range(len(arr)):
        if arr[i] == s1:
            pos1 = i
            if pos2 != -1:
                ans = min(ans, i-pos2)
        elif arr[i] == s2:
            pos2 = i
            if pos1 != -1:
                ans = min(ans, i-pos1)

    return -1 if ans == sys.maxsize else ans

进阶

如果需要查询非常多次，如何做到O(1)查询？

思路

记录每对s1, s2的最近距离。

实现

class Record():
    def __init__(self):
        self.record = {}

    def update(self, index_map, s, index):
        if s not in self.record:
            self.record[s] = {}

        s_map = self.record[s]
        for k, i in index_map.items():
            if k == s:
                continue

            k_map = self.record[k]
            cur_min = index - i
            if k in s_map:
                pre_min = s_map[k]
                if cur_min < pre_min:
                    k_map[s] = cur_min
                    s_map[k] = cur_min
            else:
                k_map[s] = cur_min
                s_map[k] = cur_min

    def setup(self, arr):
        index_map = {}
        for index, s in enumerate(arr):
            self.update(index_map, s, index)
            index_map[s] = index

    def min_distance(self, s1, s2):
        if s1 is None or s2 is None:
            return -1

        if s1 == s2:
            return 0

        if s1 in self.record and s2 in self.record[s1]:
            return self.record[s1][s2]

        return -1

测试

def test_min_distance():
    arr = []
    for _ in range(20):
        size = random.randint(0, 10)
        s = ''.join([random.choice(string.ascii_letters) for _ in range(size)])
        arr.append(s)

    arr *= 5
    random.shuffle(arr)

    record = Record()
    record.setup(arr)

    for i in range(len(arr)):
        for j in range(len(arr)):
            s1 = arr[i]
            s2 = arr[j]
            d1 = min_distance(arr, s1, s2)
            d2 = record.min_distance(s1, s2)
            if d1 != d2:
                print(s1, s2, d1, d2)
                raise Exception('Error')

    print('done')


if __name__ == '__main__':
    test_min_distance()

测试2

牛客网
python代码超时。。。C++写了一遍。。。