記一次算法面試題

前幾天收到了某所AI算法崗的面試邀請，被要求24小時內做完4道編程題，可以上網搜、找人幫忙，提交可執行代碼和解題思路。除了參加比賽和當年研究生複試，就再也沒被要求過寫算法代碼了，好吧反正解題也好玩，就寫唄。晚上九點纔拿到題，大致掃了一眼，前三題都不難，第一題考察位運算、第二題貪心算法、第三題類似動態規劃也不難，不到12點就搞定了外加解題文檔，第四題是按規則生成給定的序列，看着也還行，需要用回溯算法解決，但是一旦考慮細節就發現有坑了，最後到第二天上午才全部搞定。只記錄一下第四題吧。
記序列：1 4 1 5 6 7 4 2 3 5 2 6 3 7爲S[7]
序列規則如下：數字n與n之間間隔n個數字，即：

• 1與1之間間隔1個數字；
• 3與3之間間隔1個數字；
• 7與7之間間隔1個數字；
  以此類推。求S[12]
• 規則很簡單，但是生成的時候就會發現，s[i]到底應該填幾合適呢？既然不知道那就只能猜了，先隨便填一個按規則生成下去，當沒法繼續填數字的時候，就證明前面填錯了，回溯算法呼之欲出咯。整體指導思想有了，怎麼寫代碼呢？最初的想法很樸素，就直接向一維數組s[]裏依次填數字啊，但是在回溯的時候發現不太容易實現，當前要填的數字即和該位置之前的狀態有關，又和其餘的位置數字有關，有點複雜不想繼續糾纏下去了太痛苦。既然這樣，就還是用空間換時間吧，設有二維數組A[2*n+1, n+1]，行和列都加了1是因爲想統一下標和數值，就浪費一行一列吧。行數爲S的個數，列數爲N，A的取值爲0或1。如果A[i,j]=1，則表明s[i]=j，這樣定義後，矩陣A每列至多兩個1，每行一個1,且有A[i, j]=A[i + j + 1, j] = 1。每次回溯的時候，將該列置0即可，然後找個堆棧記錄上次填充的是哪一行，便於回溯。源碼：

    def friend_seq(N):
        # 爲了統一下標和數值，捨棄第一行和第一列不使用
        # 矩陣work_matrix行數索引代表序列S的下標，列索引代表S在此處的取值
        # 矩陣work_matrix每行只有一個值爲1，即A[i,j]=1，有：序列s[i]=j
        # 根據規則生成時有A[i, j]=A[i + j + 1, j] = 1，i,j都從1開始計數
        work_matrix = np.zeros((2 * N + 1, N + 1), dtype=np.int)
        work_matrix[2 * N, N] = work_matrix[2 * N - N - 1, N] = 1
        col_sum_vec = np.ones(N) * 2
        i = j = 1
        fill_row_index = []
        while i < 2 * N + 1:
            if np.sum(work_matrix[i, :]) == 0:
                while j < N + 1:
                    if i + j + 1 < 2 * N + 1 and np.sum(work_matrix[i + j + 1, :]) == 0 and np.sum(work_matrix[:, j]) < 2:
                        work_matrix[i, j] = work_matrix[i + j + 1, j] = 1
                        fill_row_index.append(i)
                        j = np.argmin(np.sum(work_matrix[1:, 1:], axis=0)) + 1
                        break
                    else:
                        j += 1
                if j == N + 1 and np.sum(np.sum(work_matrix[1:, 1:], axis=0) - col_sum_vec) != 0:
                    # 列已經遍歷完，沒有可用項，回溯
                    temp_i = fill_row_index[-1]
                    temp_j = np.argmax(work_matrix[temp_i, :])
                    # 更新i和j
                    i = np.argmax(work_matrix[:, temp_j])
                    j = temp_j + 1
                    # 將上一次的賦值清0
                    work_matrix[:, temp_j] = 0
                    fill_row_index.pop()
                else:
                    i += 1
            else:
                i += 1

        s = np.argmax(work_matrix[1:, 1:], axis=1) + 1
        print('{0} friend sequence: {1}'.format(N, s))

最後生成的序列：

7 friend sequence:  [1 4 1 5 6 7 4 2 3 5 2 6 3 7]
12 friend sequence: [ 1  2  1  3  2  8  9  3 10 11 12  5  7  4  8  6  9  5  4 10  7 11  6 12]

後記：將矩陣的最後一列初始化

work_matrix[2 * N, N] = work_matrix[2 * N - N - 1, N] = 1

可以提高大概10倍的效率，相當於做了一次剪枝。我猜測可能不是所有的數字N都有滿足以上規則的序列，沒測試過，還得繼續搬磚啊！