# AP325-Python 第2章 排序與二分搜
「完全無序,沒有效率;排序完整,增刪折騰;完美和完整不是同一回事。」
(sorted array vs. balanced tree)
「上窮碧落下黃泉,姐在何處尋不見,人間測得上或下,不入地獄,就是上天;
天上地下千里遠,每次遞迴皆減半,歷經十世終不悔,除非無姐,終能相見。」
(binary search)
本章介紹排序與二分搜的基本用法與應用,同時也介紹其推廣運用,包括快速冪與折半枚舉。
---
## 排序
排序是將一群資料根據某種順序由小到大排列,最常見的順序就是數字的由小到大或是由大到小,當然也可能是字元字串或是其他使用者自訂的順序。排序是非常重要的程序,計算機科學在發展初期就研發了很多種的排序演算法,這些演算法目前依然常常用在教學上當作例子與學習教材。在這裡我們主要介紹它的應用技巧而非排序演算法,對排序演算法有興趣的人可以查網路。
排序通常有幾個運用的時機:
* 需要把相同資料排在一起
* 便於快速搜尋
* 做為其他演算法的執行順序,例如Sweeping-line, Greedy, DP。
在考試與競賽的場合,最重要的是會使用庫存的排序函數,在Python是sort()與sorted(),它們非常的有效率,如果要自己寫出這麼有效率的排序需要花費很多的時間還不一定能寫好。以下我們先介紹它的使用方式與考試競賽的使用技巧。
### Python排序函數的使用方法
Python有兩個排序的函數,list.sort()與sorted(iterable)。
如果將一個list本身排序,我們使用list.sort(),執行後list的內容會被排序;如果我們要得到一個排序後的結果放在一個list中傳回,我們使用sorted(iterable),其中iteable是任何可以疊代的物件,iterable可能是list, string, set等等。這裡稍微注意list.sort()是不會回傳東西的,而sorted(iterable)是會回傳排序好的list。以下是一些要注意與容易犯錯的地方。
``` python
p = [5,3,1,4]
q = p.sort() # 是錯誤的寫法
q = sorted(p) # 是可以的
# 若要將一個字串的字元排序
s = 'qwerty'
s.sort() # 是錯誤的,因為字串沒有sort函數也無法排序
# 可以寫
q = sorted(s) # 會產生字元排序後的list
>>> q
['e', 'q', 'r', 't', 'w', 'y']
# 也常常這樣用
>>> for c in sorted(s):
print(c, end=',')
e,q,r,t,w,y,
```
這兩個函數的叫用方式十分類似,如果還不熟悉,可以先看看[線上文件](https://docs.python.org/3/howto/sorting.html)。以下我們介紹解題時常用到的使用方法與注意事項。
#### 遞增遞減與stable
Python的排序預設是遞增,如果要改成遞減,可以給關鍵字參數reverse=True。例如
```python
>>> p = [3,2,5,1,7]
>>> sorted(p)
[1, 2, 3, 5, 7]
>>> p.sort(reverse=True)
>>> p
[7, 5, 3, 2, 1]
>>>
```
Python的sort是**stable sort**,意思是對於key值相同的元素,排序後不會改變原來的前後關係。例如
```python
>>> p = [(5,7),(2,4),(5,1),(2,8),(5,4)]
>>> sorted(p, key=lambda x:x[0])
[(2, 4), (2, 8), (5, 7), (5, 1), (5, 4)]
>>>
```
此例中p是tuple的list,我們排序時要求key是看tuple的第0個欄位,所以(5,7), (5,1), (5,4)三個的key值是一樣的,所以會保持原來的前後關係。
#### 資料的大小關係
Python在資料比較時,有預設的大小關係定義,這是排序的主要依據,其實也用在max/min之類的函數。
1. 數字當然數值小就是小。
2. 字元的關係是以ASCII code的大小來定義,常見的數字字母的ASCII code
'0'<'1'<...<'9' <'A'<'B'...<'Z' < 'a'<'b'...<'z'
3. 兩個字串相比是依據字典順序(lexicographic order),由前往後找到第一個相異的來比大小。也就是說,第一個先比,如果第一個相同,再比第二個字元,如果第二個也相同就再比第三個字元,...。例如'aaa'<'b','aba'<'ac','abc'<'abd'。如果比到某個字串全部都結束還是相同,也就是一個字串是另外一個的prefix,那麼,短的比較小。
4. 兩個tuple或兩個list比較時,也是lexicographic order。但個別元素之間必須是可以比較的,數字對數字,字元對字元,字串對字串tuple對tuple,list對list。
#### key function
sort()與sorted()內的重要參數是key。前面所說的大小關係是預設的,如果要改變,我們必須透過key這個參數來指定一個函數,做為排序過程中取得key值的依據。舉例來說,我們要排序一些區間,每個區間是(left,right)表示左右端點。如果不指定key,那就是字典序,也就是依照左端點由小到大,左端點相同的會比較右端點。如果我們想用右端點來排序,可以像下面這樣:
```python
>>> def get_right(p):
return p[1]
>>> w = [(3,5),(1, 6),(4,2)]
>>> sorted(w,key=get_right)
[(4, 2), (3, 5), (1, 6)]
>>>
```
像這樣特別定義一個函數來給sort叫用,但在別處並不使用時,會覺得很麻煩,所以通常會使用**無名函數(lambda function)** 的方式,其中Lambda是希臘字母$\lambda$。也就是說,我們只用在此處,所以省略不取名字的函數,做法是在key的地方直接給定義,像下面這樣。
```python
>>> sorted(w, key=lambda x:x[1])
[(4, 2), (3, 5), (1, 6)]
```
lambda x: x[1] 代表傳入的參數如果叫做x,key就取x[1]。
如果我們想用兩個欄位之和來做為key,可以這樣寫
```python
>>> sorted(w, key=lambda x:x[0]+x[1])
[(4, 2), (1, 6), (3, 5)]
```
我們在下面的範例題目中會看到很多排序的例子。另外,在有些狀況下,key function不好定義,而用一個比較函數來做,再透過functools中的cmp_to_key轉成key函數。
**基本範例**
下面的例題是排序的一個應用,常用在某些題目的前置處理,為了方便練習,我們把它寫成題目的樣子。
---
##### P-2-1. 不同的數 --- 排序
假設有N個整數要被讀到一個陣列中,我們想要將這些數字排序並去除重複的數字,例如輸入的整數序列是 (5, 3, 9, 3, 15, 9, 8, 9),這些數如從小到大排是 ( 3, 3, 5, 8, 9, 9, 9, 15),去除重複者後為(3, 5, 8, 9, 15)。寫一個函數,傳回有多少不同的數字並且將結果放在陣列中回傳。
Time limit: 1秒
輸入格式:輸入兩行,第一行是正整數$N$,$N$不超過10萬,第二行是$N$個整數,大小不超過$10^9$,以空白間隔。
輸出: 第一行輸出有多少相異整數,第二行由小到大輸出這些相異整數,相鄰數字之間以一個空白間隔。
範例輸入:
7
0 3 9 3 3 -1 0
範例結果:
4
-1 0 3 9
---
請看以下的範例程式。函數distinct()傳入整數陣列。我們不希望破壞原陣列的內容,將資料排序後放在另外一個陣列 $v$,排序後相同的數字會在一起。初始一個空列表$res$來放置最後的結果,我們先將$v[0]$放入,然後歷遍$v$的其餘元素,只要$v[i]$與$res[-1]$(最後一個)不同的話,就是一個相異的數字。
```python=
# P_2_1 distinct number -- sort
def distinct(arr): # return distinct number in arr
if len(arr) == 0: return res
v = sorted(arr)
res = [v[0]]
for p in v[1:]:
if p != res[-1]: # distinct
res.append(p)
return res
#
n = int(input())
a = [int(x) for x in input().split()]
p = distinct(a)
print(len(p))
print(*p)
```
註:如果運用set(),distinct()可以用一行實作
```python=
def distinct(arr): # return distinct number in arr
return sorted(set(arr))
```
其中先將$arr$放入一個集合中,集合只會保留相同的元素,然後回傳將集合元素奇排序後的結果,排序後會轉成一個list。如果是初學的讀者,還是要學會一些基本的集合操作比較好。
## 搜尋
在一群資料中搜尋某筆資料,通常有線性搜尋與二分搜尋,線性搜尋就是一個一個找過去,在沒有順序的資料中,也只好如此。如果資料已經依照順序排好,我們可以採用二分搜尋,原因是效率好太多了。以一個100萬筆的資料來說,線性搜尋可能需要比對100萬次,而二分搜尋的worst case複雜度是$O(log(n))$,最多只要比較20次就可以找到了,也就是有5萬倍的效率差異,所以對很多問題是非用不可。
二分搜除了自己寫之外,也有庫存函數可以呼叫,此外,C++也有一些好用的資料結構可以達到搜尋的目的,我們在這一節中會介紹以下技巧:
* 基本二分搜的寫法
* Python的二分搜函數
此外,二分搜尋不單是在一群資料中找資料,它往往也搭配其他演算法策略使用,這一部份會出現在後續的章節中。
### 基本二分搜的寫法
最簡單的二分搜是在一個排好序的陣列中尋找某個元素,找到了回傳index,否則回傳找不到(例如-1),這樣的二分搜好寫也不容易寫錯,以下是個範例。二分搜的重點在於left與right兩變數紀錄著搜尋範圍的左右邊界,每次取出中間位置來比較,結果是找到了或捨棄左半邊或捨棄右半邊。當left>right表示搜尋區間已經沒有了,這時離開迴圈,因為要找的元素不存在。若初始的區間範圍是$n$,時間複雜度是$O(log(n))$,原因是區間長度每次都減半。
```python=
import random
# binary search x between a[left..right]
def bsearch(a, left, right, x):
while left <= right:
mid = (left+right)//2 # middle element
if a[mid] == x: return mid
if a[mid] < x: left = mid+1 # search right part
else: right = mid-1# search left part
#
return -1
#
p = [random.randrange(100) for i in range(10)]
n = len(p)
p.sort()
print(*p)
print("search", p[7], "=> return",bsearch(p,0,n-1,p[7]))
t = random.randrange(100)
print("search", t, "=> return",bsearch(p,0,n-1,t))
```
在很多時候,我們需要知道不只是某元素在不在,當它不在時,希望找到小於它的最大值的位置(或是大於它的最小值),這看似只要把上述程式簡單修改一下就好了,但事實上很容易寫錯。假設我們需要知道第一個大於等於$x$的位置,我們將前述程式直接修改如下。
```python=
# find the first >=x between a[left..right]
def bsearch(a, left, right, x):
while left <= right:
mid = (left+right)//2 # middle element
if a[mid] == x: return mid
if a[mid] < x: left = mid+1 # search right part
else: right = mid-1# search left part
#
return ??? # maybe wrong
```
你檢查後可能發現應該要回傳left是對的,事實上很多人會弄錯,另外也會不小心寫成無窮迴圈。事實上,上面那支程式還是有錯誤,在有相同元素時,不能找到一個要搜尋的元素就回傳,因為前面可能也有相同元素。
重新審視這個程式,我們維護以下迴圈不變性:要找的位置位於$[left,right]$這個閉區間內。
```python=
# find the first >=x between a[left..right]
def bsearch(a, left, right, x):
if a[right] < x: return -1
while left < right: # more than 1
mid = (left+right)//2 # middle element
if a[mid] < x: left = mid+1 # search right part
else: right = mid# search left part
#
return left
```
一開始,我們必須先確定迴圈不變性是滿足的(第3行),在此不變性之下,如果只剩下一個元素就是找到了,所以迴圈繼續搜尋的條件是區間長度大於1 (第4行),跟中間位置比對後,如果$a[mid] <x$,$\le mid$就不是要搜尋的對象,所以應該把$left=mid+1$;反之,$>mid$不是要搜尋的對象,所以$right=mid$。
當迴圈結束時,根據迴圈不變性,答案必在區間內,此時$left==right$就是答案。
維護迴圈不變性是一個很重要的思考方式,它也是程式正確性的證明。
二分搜有很多不同的場景,這個區間搜尋的寫法並非一定不好,這裡只是希望提醒讀者這個寫法容易犯錯。筆者比較推薦的寫法在前一章曾經提出過,不是用左右搜的方式,而是用一路往前跳的方式,你想想看,正常人如果反覆的向左向右是不是很容易昏頭,一路往前應該比較不容易搞錯。
請看以下的範例程式,要在一個list $a$中找到第一個$\ge x$的位置,如果不存在則回傳$a$的長度。
我們用$po$作為目前的位置,在一開始檢查一下$a[0]$是否就是所要答案,如果不是,程式中就會一直保持$a[po]<x$這個不變性,在此特性下,只要還能往前跳就往跳,跳的距離每次把它折半。
```python=
# binary search the first >=x between a[0..n-1]
def jump_search(a, x):
if a[0] >= x: return 0 # check the first
po = 0 # current position, always < x
jump = len(a)//2 # initial jump distance
while jump > 0:
while po+jump <len(a) and a[po+jump]<x:
po += jump
jump //= 2
return po+1
```
時間複雜度也是$O(log(n))$,因為跳的距離每次折半,而且內層的while迴圈至多執行兩次(不會很難證明)。
### Python的二分搜函數
Python提供了二分搜模組許多跟搜尋有關的函數以及資料結構,在考試或比賽中,如果是適合的題目,使用庫存函數可以節省時間也避免錯誤,這裡介紹一些基本的用法。Python的二分搜在bisect模組內,完整的說明可以參考 [官方文件(3.6版)](https://docs.python.org/zh-tw/3.6/library/bisect.html),使用前需要
import bisect
最基本需要也是最常用的函數有兩個,其基本使用方式如下:
* bisect.bisect_left(a,x):在遞增的列表a中找到第一個 $\ge x$的位置(index),如果不存在則回傳陣列長度,也就是最後一個元素的下一個位置。
* bisect.bisect_right(a,x):與前者類似,但找到的是第一個$>x$的位置。
時間複雜度都是$O(\log n)$。這兩個函數都還可以給搜尋的左端與右端的參數,大部分場合不太需要,預設就是搜尋整個列表。
註:新版的還可以給與key的參數。
以下這個範例程式很簡單,可以把它拿來執行一下就可以了解bisect的基本運用。
```python=
# demo bisec
import bisect
p = sorted([5, 1, 8, 3, 9])
print(p)
for t in range(0,15,3):
idx = bisect.bisect_left(p, t)
if idx < len(p):
print("The first >=",t,"is at", idx)
else: print("No one >=",t)
idx = bisect.bisect_right(p, t)
if idx < len(p):
print("The first >",t,"is at", idx)
else: print("No one >",t)
```
執行結果:
```
[1, 3, 5, 8, 9]
The first >= 0 is at 0
The first > 0 is at 0
The first >= 3 is at 1
The first > 3 is at 2
The first >= 6 is at 3
The first > 6 is at 3
The first >= 9 is at 4
No one > 9
No one >= 12
No one > 12
```
當列表中的元素也是列表(或元組)的型態,例如兩個欄位的平面座標或是多欄位的資料,我們還是可以使用bisect,記得比較大小的方法是lexicographic order,以下是個簡單的範例程式。
```python=
import bisect
p = [[3,1],[1,4],[3,5],[4,2],[0,3]]
p.sort()
print(p)
n = 5
t = [3,2]
idx = bisect.bisect_left(p, t)
print("Find >=",t, "at", idx)
t = [4]
idx = bisect.bisect_left(p, t)
print("Find >=",t, "at", idx)
```
執行結果:
```
[[0, 3], [1, 4], [3, 1], [3, 5], [4, 2]]
Find >= [3, 2] at 3
Find >= [4] at 4
```
下面一個例題是一個有趣的排序與搜尋練習,也是實際解題時經常需要用的步驟,它有個名字是「離散化」,雖然這個名字似乎不是那麼恰當(英文名字是座標壓縮coordinate compression),但是從以前就這麼多人這樣稱呼它。它需要用到前一個例題P_2_1(不同的數)當作它的步驟。
---
##### P-2-2. 離散化 --- sort
假設有$N$個整數要被讀到一個陣列中,我們想要將這些整數置換成從0開始的連續整數並且維持它們原來的大小關係,例如輸入的整數序列是 (5, 3, 9, 3, 15, 9, 8, 9),這些數如從小到大排是 ( 3, 3, 5, 8, 9, 9, 9, 15),去除重複者後為(3, 5, 8, 9, 15),所以我們要替換的是:
3 -> 0
5 -> 1
8 -> 2
9 -> 3
15 -> 4
所以原先的序列就會變成(1, 0, 3, 0, 4, 3, 2, 3)。
Time limit: 1秒
輸入格式:輸入兩行,第一行是正整數$N$,$N$不超過10萬,第二行是$N$個整數,大小不超過$10^9$,以空白間隔。
輸出:輸出置換後的序列,兩數之間以一個空白間隔。
範例輸入:
7
0 3 9 3 3 -1 0
範例輸出:
1 2 3 2 2 0 1
---
如果已經了解如何處理不同的數字,加上二分搜尋就可以簡單地處理這個問題,請看以下範例。我們先呼叫剛才寫好的函數將陣列中的相異數字排序在b[]陣列中,然後,對於原陣列a[]中的每一個數字,將其置換成它在b[]中的index就可以了,而要快速的在b[]中找到a[i],我們用二分搜,自己寫或是呼叫bisect_left()都可以,因為a[i]一定在b[]中,所以我們不需要擔心找不到的情形。
```python=
# P_2_2 discretization -- sort
import bisect
def distinct(arr): # return distinct number in arr
if len(arr) == 0: return res
v = sorted(arr)
res = [v[0]]
for p in v[1:]:
if p != res[-1]: # distinct
res.append(p)
return res
#
n = int(input())
a = [int(x) for x in input().split()]
b = distinct(a) # find and sort distinct
for i in range(n): # replace with rank
a[i] = bisect.bisect_left(b,a[i])
print(*a)
```
如果使用set與dict來做就更簡單了,以下是範例程式,如果暫時沒學set與dict,可以先跳過這個寫法。
```python=
# P_2_2 discretization -- set and dict
n = int(input())
a = [int(x) for x in input().split()]
# build dict mapping to rank
d = {p:i for i,p in enumerate(sorted(set(a)))}
for i in range(n): # replace with rank
a[i] = d[a[i]]
print(*a)
```
### 其他相關技巧介紹
在這一節中,我們將介紹以下主題:
* Bitonic sequence的搜尋
* 快速冪
#### Bitonic sequence的搜尋
Bitonic sequence是指先遞增後遞減(或者先遞減再遞增)的序列,也有人叫它一山峰(一山谷)序列。在單調遞增序列(函數)中可以用二分搜來找第一個超越0值或某值的位置,那麼對於Bitonic序列是否可以類似二分搜有效率地找到極值呢?例如在一山谷的序列中找最小值。答案是有的,常用的方法有兩種:三分搜以及差分二分搜。
與二分搜一樣,我們始終維護一個搜尋區間,二分搜是每次將區間二等分,一個比較之後刪除一半。三分搜是將區間三等分,如果1/3的位置大於2/3的位置,可以丟掉左邊的1/3;反之,如果2/3的位置較大,則可以丟棄右邊1/3的區間。如此每次可以將區間長度減少1/3,在$O(\log(n))$次比較可以找到最低點。
另外一個方法更簡單,若序列$f$為先遞減再遞增(一山谷),則差分
$g[i]=f[i]-f[i-1]$
在前半段是負的,在後半段是正的,只要以二分搜找到差分序列第一個大於等於0的位置就是$f$的最小值。
#### 快速冪
所謂快速冪是如何快速計算$x^y$的方法,這裡要探討的不是浮點數的運算,通常是針對整數,而這個數字通常都大到超過整數變數的範圍,所以通常都是求模(mod)$P$的運算,也就是給定$x$, $y$, $P$,要計算$x^y$除以$P$的餘數。
當數字比較小的時候,可以用$x**y$計算次方,但數字大的時後不可以,不過,Python的內建函數$pow()$已經是以快速冪的方式計算$x^y$除以$P$的餘數,甚至$pow(x,-1,P)$可以直接計算乘法反元素(模逆元)。但是下面要講的方法還是需要學,因為在某些場合,我們還是要自己寫快速冪,而且這也是一種很重要的基本技巧。
---
##### P-2-3. 快速冪
輸入正整數$x$, $y$, 與$p$,計算 $x^y$ (mod $p$),其中$x$, $y$, $p$ 皆不超過1e9+9。例如$x=2$, $y=5$, $p=11$,則答案是10。
Time limit: 1秒
輸入格式:輸入$x$, $y$, 與$p$在同一行,以空白間隔。
輸出格式:輸出計算結果。
範例輸入:
2 5 11
範例輸出:
10
---
計算指數最直接的方法就是按照定義一個一個的乘,如下列的程式碼,提醒一下,每次運算後都要取餘數,否則有overflow 或是運算太慢的可能。
```python
t = 1
for i in range(y):
t = t*x%p
```
但是這個方法效率太差,如果y是10億就要執行10億次乘法運算。前面介紹二分搜逐步縮減跳躍距離的方法,以類似的概念,我們可以設計出只要$\log(y)$次乘法的方法,這個方法以遞迴的形式最簡單。在以下範例中我們同時展示遞迴與迴圈兩種寫法。
```python=
# P_2_3 find x^y mod P, recursion and loop
def exp(x, y, p): # recursive
if y==0: return 1
if (y & 1): return exp(x, y-1,p)*x%p
# otherwise y is even
t = exp(x, y//2, p)
return t*t%p
#
def exp2(x, y, p): # loop
t = 1; xi = x; i = 1 # t is result, xi = x^ (2^i)
while y > 0:
if (y & 1): # odd, (i-1)-bit of y = 1
t = t*xi%p
y >>= 1
xi = xi*xi%p
i = i*2 # i is useless, just for explanation
return t
#
x, y, p = [int(x) for x in input().split()]
res = exp(x, y,p)
print(res)
# check two answer
if res != exp2(x,y,p):
print("different result")
```
遞迴的寫法很好懂,若$y$是奇數則先遞迴求出$y-1$次方後再乘一次$x$;若$y$是偶數則求出$y//2$次方後自乘,終止條件為$y=0$時結果是$1$。遞迴的版本中,$xi$是每次自乘的結果,所以它的內容是$x$的1次方、2次方、4次方、8次方、…這樣的方式變化,將$y$以二進制來看,對每一個bit為1的位置,把對應項的$x$次方乘到結果中就是答案。舉例來說,若$y=19$,二進制是10011,就是要計算
$x^{19} = x^{16} \times x^2 \times x^1$。
兩種寫法的效率差不多,都是很快的方法,遞迴版本尤其容易記。
以下是一個稍加變化的習題,給一個提示:當$x$很大很大時,我們可以先求$x$除以$p$的餘數,而除以$p$的餘數可以一位一位的累加計算。例如$x=123$, $p=5$,$x$可以看成$(1\times 10+2)\times 10+3$,所以$x$除以$p$的餘數也就等於
$((((1\times 10\%p+2)\%p)\times 10)\%p+3)%p$。
---
##### Q-2-4. 快速冪--5000位整數
題目與輸入皆同於P-2-3,但$x$的範圍是不超過5000位的正整數
範例輸入:
123456789012345678901234567890 5 11
範例輸出:
10
註:新版本的Python可以透過設定增加整數位數的長度,如此可以直接的計算$x$除以$p$的餘數,但前面提示的方法還是值得學習的。
---
**快速計算費式數列**
利用快速冪我們可以計算費式數列的第$n$項。
---
##### Q-2-5. 快速計算費式數列第n項
令$f[0]=0$, $f[1]=1$, 以及 $f[n]=f[n-1]+f[n-2]$ for $n>1$。輸入非負整數$n$,請輸出$f[n]$除以$p$的餘數, $p=1000000007$。$n<2^{31}$。
Time limit: 1秒
輸入格式:輸入可能有多行,每一行有一個整數是一筆測資,最後一行以-1代表結束,不需要處理該筆測資。
輸出格式:每一行依序輸出計算結果。
範例輸入:
6
123456789
100
-1
範例輸出:
8
62791945
687995182
---
Q_2_5說明:我們可以將費式序列的定義寫成以下矩陣的形式
$$
\left(
\begin{array}{c}
f[n]\\
f[n-1]
\end{array}\right)
= \left(
\begin{array}{cc}
1 & 1\\
1 & 0
\end{array}\right) \times
\left(
\begin{array}{c}
f[n-1]\\
f[n-2]
\end{array}\right)
$$
將上式迭代展開,因為矩陣滿足結合律,所以可以得到
$$
\left(
\begin{array}{c}
f[n]\\
f[n-1]
\end{array}\right)
= \left(
\begin{array}{cc}
1 & 1\\
1 & 0
\end{array}\right) ^{n-1} \times
\left(
\begin{array}{c}
f[1]\\
f[0]
\end{array}\right)
$$
如果我們可以很快的算出那個矩陣$A$的$n-1$次方,就可以得出
$f[n]=A^{n-1}[0][0]$。
## 2.4. 其他例題與習題
排序與搜尋大多都會與其他的演算法結合,成為某個解題方法中的一個步驟,單獨考排序的並不多。以下舉一些例子,先看一個很常見的基本問題。
---
##### P-2-6. Two-Number problem
假設$A$為$m$個相異整數的集合,$B$為$n$個相異整數的集合,而$K$是一個整數。請計算有多少對$(a, b)$的組合滿足$a \in A$, $b\in B$ 且$a+b = K$。
Time limit: 1秒
輸入格式:輸入可能有多行,第一行有三個整數$m$, $n$與$K$,第二行有$m$個整數是$A$中的元素,第三行有$n$個整數$B$中的元素。同一行相鄰數字間以空白間隔。兩集合元素個數均不超過10萬,整數的絕對值不超過10億。
輸出格式:輸出組合個數。
範例輸入:
3 4 2
1 6 -3
5 1 -1 -3
範例輸出:
2
---
這一題有多個解法,基本的想法是排序後搜尋,有下列幾種作法:
* 將$A$排序後,對$B$中的每一個$b$,以二分搜在$A$中尋找$K-b$。
* 將$A$與$B$分別排序後,對$A$中的每一個$a$,在B中以滑動的方式找$K-a$。
* 把$A$放進set來做
由於作法類似,我們省略過第一種。以下是第二種的範例程式,在$a$與$b$分別排序後,由前往後對於每一個$a[i]$,以while迴圈由後往前找到$b$中第一個小於等於 $k - a[i]$的位置,請注意,由於$a[i]$是由小到大,因此$j$每次不必從最尾端重新開始,只要從前一次停下的位置開始就可以了,所以如果不計排序,這個程式的複雜度是$O(m+n)$。這個在陣列中維護兩個位置的方法,也有人稱為Two-pointers method(雙指針方法),它雖然用了雙迴圈但複雜度並不是平方,類似的情形我們在下一章會看到更多例子。
```python=
# p_2_6a find a+b=k
m, n, k = [int(x) for x in input().split()]
a = [int(x) for x in input().split()]
b = [int(x) for x in input().split()]
a.sort() # sort a from small to large
b.sort() # sort b from small to large
j = n-1 # index of b, from n-1 to 0
ans = 0
for ai in a: # each a[i]
while j>0 and b[j]>k-ai: # backward linear search
j -= 1
if ai + b[j] == k: ans +=1
#
print(ans)
```
接下來看用利用set的做法,請看以下範例程式,我們可以注意到,在這個例子上我們甚至不需要用到任何陣列。
```python=
# p_2_6b find a+b=k, using set
m, n, k = [int(x) for x in input().split()]
s = set([int(x) for x in input().split()])
ans = 0
for bi in [int(x) for x in input().split()]:
if k-bi in s: # O(1) hash search k-bi
ans += 1
#
print(ans)
```
接下來是一個與前者類似但較複雜一點的習題。
---
##### Q-2-7. 互補團隊 (APCS201906)
前 $m$ 個英文大寫字母每個代表一個人物,以一個字串表示一個團隊,字串由前 $m$ 個英文大寫字母組成,不計順序也不管是否重複出現,有出現的字母表示該人物出現在團隊中。兩個團隊沒有相同的成員而且聯集起來是所有$m$ 個人物,則這兩個團隊稱為「互補團隊」。輸入 $m$ 以及$n$ 個團隊,請計算有幾對是互補團隊。我們假設沒有兩個相同的團隊。
Time limit: 3秒
輸入格式:第一行是兩個整數 $m$ 與 $n,$2 \le m \le 26$,$1 \le n \le 50000$。第二行開始有 $n$ 行,每行一個字串代表一個團隊,每個字串的長度不超過$100$。
輸出格式:輸出有多少對互補團隊。
範例輸入:
10 5
AJBA
HCEFGGC
BIJDAIJ
EFCDHGI
HCEFGA
範例輸出:
2
---
解題提示:與前面的例題在結構上是相似的,前面是找數字,這裡是找集合。因為字串中有重複的字母而且沒有照順序,我們需要將每一個集合表示成唯一的表式方式才能夠有效的搜尋。有兩個方法可以做:
* 將每個字串去除重複字母且裡面的字母是由小到大排列的。例如AJBA就改成ABJ,HCEFGGC就改成CEFGH。
* 以一個整數表示一個集合,第$i$個bit設為1代表第$i$個字母在集合中,否則為0。此法在在第一章窮舉子集合時也介紹過。
第二種方法要比第一種方法更快,因為數字的搜尋要比字串來得快。以下的程式片段可以將一個字串轉換成對應以及互補集合的整數:
```python
#transforming a string to corresponding int
ff = (1<<m) - 1 # bits 0~(m-1) are 1
s = input()
team = 0
for ch in s: # each char of s
team |= 1<<(ord(ch) - 0x41) # set bit to 1
# ord(ch) is ASCII of ch, 0x41 = ord('A')
complement = ff - team # int of the complement
```
以下的習題我們用來展示如何求Modular multiplicative inverse,就是在模運算下的乘法反元素,有些人稱為「模反元素」或「模逆元」。
對於任何一個正整數$a$,它的模$P$乘法反元素就是滿足$(a \times b) \% P = 1$ 的整數$b$,這裡的 \% 就是取餘數運算。計算$a$的模逆元是一個很重要的運算也有許多運用,最簡單的方法是如以下的窮舉測試:
```python
for b in range(1,P):
if (a * b) % P == 1:
# b is an inverse
```
窮舉法的問題在於效率太差。在許多運用場合$P$是一個質數,而且探討的整數範圍都只在0 ~ $P-1$。在這些假設下有一個重要的數學性質可以幫助我們快速的計算模逆元(費馬小定理):「若$P$為質數,對任意正整數$a$,$(a^{P-2}\ \% P)$是$a$在$[1, P-1]$區間的唯一乘法反元素。」根據此性質搭配快速冪就可以用$O(\log P)$個運算計算出模逆元。另一種求模逆元的方法是用Extended Euclidean algorithm,這裡就不介紹了。
---
##### Q-2-8. 模逆元 (*)
輸入$n$ 個正整數,以及一個質數$P$,請計算每一個輸入數的模逆元。輸入的正整數的大小不超過$P$,$P\le 1000000009$,$0 < n < 10$。
Time limit: 1秒
輸入格式:第一行是$n$與$P$,第二行$n$個整數,同行數字以空白間隔。
輸出格式:依照輸入順序輸出每一個數的模逆元,相鄰數字間間隔一個空白。
範例輸入:
3 7
3 4 1
範例輸出:
5 2 1
註1:費馬小定理屬於數論中的定理,應該不在APCS考試的背景知識內,但如果題目中給予提示說明,則模逆元並非屬於一定不可以考的範圍。
註2:Python 在3.8版之後,內建函數pow可以直接計算模逆元,但之前版本的Python有快速冪但無法直接計算負的次方,負的次方會轉成float。
---
透過下面的例題,我們要來介紹折半枚舉的技巧,這一題算是有難度的題目。
---
##### P-2-9. 子集合乘積(折半枚舉) (@@)
輸入$n$個正整數$A[1..n]$,以及一個質數$P$,請計算$A$中元素各種組合中,有多少種組合其相乘積除以$P$的餘數等於1。每個元素可以選取或不選取但不可重複選,$A$中的數字可能重複。$P\le 1000000009$,$0 < n \le 32$,且假設$A$中元素皆小於$P$。
Time limit: 4秒
輸入格式:第一行是$n$與$P$,第二行$n$個整數是$A[i]$,同行數字以空白間隔。
輸出格式:滿足條件的組合數,因為數字可能太大,請輸出該組合數除以$P$的餘數。
範例輸入:
5 11
1 1 2 6 10
範例輸出:
7
說明:乘積等於1的組合有:(1), (1), (1,1), (2,6), (1,2,6), (1,2,6), (1,1,2,6)共7種。
---
我們可以透過枚舉所有的子集合來測試哪些的組合乘積為1。在第一章我們說明過以遞迴來枚舉所有子集的方法,他的效率是$O(2^n)$,在這裡$n$可能達到32,顯然無法在一秒內完成。以下的範例程式是一個修改的方法,在相同的乘積多的場合他的效率會有所改善,但不足以通過這一題所有測資。這個程式的方法也很直接,我們逐一考慮每一個元素,以一個defaultdict(int) d1來記錄著目前所有可能的子集合乘積,相同的乘積只記錄一筆,但記錄能產生此乘積的子集合個數。對於下一個元素a[i],把原有的所有可能在乘上此新元素,就是新的可能乘積。在計算過程,新的乘積不能直接加到原來的裡面,否則無法分辨新舊。因此要利用一個d2來暫存,當一個元素處理完畢後交換d1與d2,以便下一個元素時再從d1計算到d2。
註:defaultdict(int)是dict的一種子類別,它與dict相似,差別在於如果key值不存在時dict[key]將會引發錯誤,但defaultdict則會產生一個int的預設值(0)。
```python=
# p-2-9 subset product = 1 mod P, slow method
from collections import defaultdict
n, p = [int(x) for x in input().split()]
a = [int(x) for x in input().split()]
d1 = defaultdict(int) # d[product]= number
for v in a: # for each element a[i]
# compute from d1 to d2
d2 = d1.copy() # copy d1 to d2
for u in d1:
t = u*v%p
d2[t] = (d2[t] + d1[u])%p
d2[v] = (d2[v] + 1)%p # for v
d1 = d2
#
print(d1[1])
```
在這種場合我們可以用折半枚舉來做的更有效率一點。
對於n個數字,我們先將他任意均分為兩半$A$與$B$,我們要找的解(子集合乘積等於1)有三種可能:在$A$中、在$B$中、以及跨$A,B$兩端(包含兩邊的元素)。我們對$A$與$B$分別去窮舉它們的子集合乘積,可以找到在$A$中與在$B$中的解。對於跨兩邊的解,我們以下列方式計算:將其中一邊(例如$B$)的所有子集合乘積收集在某個容器中後,我們對每一個$A$的子集合乘積$x$,在$B$的子集合乘積中去搜尋$x$的模逆元$y$的個數(使得$xy = 1$的$y$)。
有一點必須留意,在$A$與$B$中都可能有很多相同元素,如果$A$中有$p$個$x$,而其模逆元$y$在$B$中有$q$個,那麼我們要避免花到$pq$的時間,因為$pq$可能會太大。解決的方法有多種,其中一個方法是,對於一個$x$搜尋時,很快地找出$y$的數量,不能很多的模逆元就會太花時間。至於$A$那一邊,合併也可以提升效率,但worst case複雜度沒有影響。我們可以算一下整個時間複雜度:對$n/2$個元素窮舉子集合需要$O(2^{n/2})$,兩邊都做所以乘以2,排序需要$O(n\times 2^{n/2})$,接著對$2^{n/2}$個數字在$2^{n/2}$個數字中做二分搜,需要的時間是$O(n\times 2^{n/2})$,所以總時間複雜度是$O(n\times 2^{n/2})$,這比起單純窮舉$O(2^n)$要快得多了。以下是範例程式,其中求模逆元是前面介紹過的快速冪方法。
```python=
# P-2-9 subset product = 1 mod P, O(n*2^(n/2)), sort
import bisect
n, p = [int(x) for x in input().split()]
a = [int(x) for x in input().split()]
# generate all products of subsets of v[]
def subset(v):
prod = []
for e in v:
prod += [q*e%p for q in prod]+[e]
return prod
# find x^y mod P
def exp(x, y, p): # recursive
if y==0: return 1
if (y & 1): return exp(x, y-1,p)*x%p
# otherwise y is even
t = exp(x, y//2, p)
return t*t%p
#
def equal(arr,x): # find number of x in arr (sorted)
return bisect.bisect_right(arr,x)-bisect.bisect_left(arr,x)
#
sa = subset(a[:n//2]) # product of half subset of a
sb = subset(a[n//2:]) # the other
sb.sort()
#print(sa,sb)
ans = equal(sb,1) # the number of 1 in sb
# compute 1 in sa and cross the two sides
for v in sa: # for each x in sa, find its inverse in sb2
if v==1: ans += 1
y = exp(v, p-2, p) # inverse
ans += equal(sb,y)
#
print(ans%p)
```
這一題當然也可以用dict來做,以下是範例程式。這裡的子集合窮舉是以遞迴方式寫的,在遞迴到最後時將所產生的子集合乘積放入dict中。
```python=
# P-2-9 subset product = 1 mod P, O(2^(n/2)), dict
from collections import defaultdict
n, p = [int(x) for x in input().split()]
a = [int(x) for x in input().split()]
# generate all products of subsets of v[], recursive
def subset(v,idx,curr,res):
if idx==len(v):
res[curr] += 1
return
subset(v,idx+1,curr,res)
subset(v,idx+1,curr*v[idx]%p,res)
return
# find x^y mod P
def exp(x, y, p): # recursive
if y==0: return 1
if (y & 1): return exp(x, y-1,p)*x%p
# otherwise y is even
t = exp(x, y//2, p)
return t*t%p
#
sa = defaultdict(int)
subset(a[:n//2],0,1,sa) # product of half subset of a
sb = defaultdict(int)
subset(a[n//2:],0,1,sb) # the other
# compute 1 in sa and cross the two sides
ans = 0
for v in sa: # for each x in sa, find its inverse in sb2
y = exp(v, p-2, p) # inverse
ans += sa[v]*sb[y]
#
print((ans-1)%p) # -empty set
```
下面我們給一個類似的習題,這個題目在上一章也出現過,當時的$n<21$,現在把他放寬到38,解法跟前面的例題很類似,注意這裡的$P$不一定也不需要是質數。因為沒有要求模逆元也不必做合併,這題比前面的例題簡單一點。
---
##### Q-2-10. 子集合的和(折半枚舉)
輸入$n$個正整數$A[1..n]$,另外給了一個整數$P$,請計算$A$中元素各種組合中,其和最接近$P$但不超過$P$的和是多少。每個元素可以選取或不選取但不可重複選,$A$中的數字可能重複。$A[i]$與$P$均不超過 $2^{60}$,$0 < n \le 37$。
Time limit: 3秒
輸入格式:第一行是$n$與$P$,第二行$n$個整數是$A[i]$,同行數字以空白間隔。
輸出格式:最接近$P$但不超過$P$的和。
範例輸入:
5 17
5 5 8 3 10
範例輸出:
16
---
在P-1-6,我們提出了一個尋找最接近區間和的問題,當時用來展示一些窮舉的技巧,只有提出$O(n^2)$的方法,以下的例題我們來展示如何有效率的解這個問題。
---
##### P-2-11. 最接近的區間和 (\*)
輸入一個整數序列$(A[1], A[2], …, A[n])$,另外給了一個整數$K$,請計算哪一個連續區段的和最接近$K$而不超過$K$。$n$不超過10萬,每個數字的絕對值不超1萬。
Time limit: 1秒
輸入格式:第一行是$n$與$K$,第二行$n$個整數是$A[i]$,同行數字以空白間隔。
輸出格式:在所有區間和中,最接近$K$但不超過$K$的和。
範例輸入:
5 10
5 -5 8 -3 4
範例輸出:
9
註:這一題是超過APCS程度的題目,因為需要用到一個比較複雜的資料結構。
---
題目並沒說輸入的是正數,這裡的輸入數字當然有可能是負的,否則有另外一種解法(下一章會討論)。一個區段可以用兩端點$[l,r]$表示,對於可能的右端$r$,我們要設法找到最好的左端$l$,最好的意思就是讓$A[l:r+1]$的和最接近$K$而不超過$K$,這裡用的是Python的表示法,$A[i:j]$是含$i$不含$j$的左閉右開區間。如此一來,然後對所有的$r$再找出最好的解就是最佳解。所以問題的核心再如何對任一右端計算出最好的解。
在P-1-6中我們介紹過前綴和(Prefix-sum),令$ps[i]$表示前$i$項的和,為方便起見定義$ps[0]=0$。對任一右端$r$,$sum(A[i:r+1])=ps[r]-ps[i-1]$,因此,要使$sum(A[i:r+1])$不超過$K$又最接近$K$,就是要找$ps[i-1]\ge ps[r]-K$的最小值。
解法呼之欲出了,讓$r$從小往大跑,以資料結構維護好所有$i<r$的$ps[i]$,然後找到第一個大於等於$ps[r]-K$的值。那麼該用甚麼資料結構呢?因為必須不斷的將prefix-sum加入,所以必須以動態的資料結構來處理,對 C\++來說,STL的set就會符合需要,時間複雜度是$O(n\log(n))$,因為在C\++的set中搜尋是$O(\log(n))$。
本題是超過APCS範圍的題目,對於Python來說,內建函數並沒有這樣的資料結構。我們以下介紹兩個解法。
第一個方法是使用二分搜與list.insert(),以下是範例程式。程式其實很簡單,用list $ps$紀錄目前所有的prefix sum並保持sorted狀態,初始放入0是代表空的前綴,$psum$紀錄目前的prefix sum,$best$紀錄目前最佳解。第7行歷遍輸入的資料,第9行用二分搜找到目標值,如果找到就嘗試更新最佳解,然後在結束這一回合前,將前綴和加入$ps$並保持排好序的狀態(第13行),作法是二分搜找到插入位置,再使用insert(),也可以用第12行bisect中提供的函數。要留意,list.insert()的時間複雜度其實是$O(n)$,所以整個程式的時間複雜度是$O(n^2)$,但程式內的主要運算都是在底層完成的函數,速度比想像中快,不過無法在時間內通過所有測資。
```python=
# p_2_11, using bisect, better O(n^2)
import bisect
n, k = [int(x) for x in input().split()]
ps = [0] # previous prefix_sum
best = -1000000000 # -oo
psum = 0 # current prefix sum
for v in (int(x) for x in input().split()):
psum += v # //prefix-sum(r)
i = bisect.bisect_left(ps, psum-k)
if i < len(ps): # found
best = max(best, psum-ps[i]) # currently best
#bisect.insort_left(ps,psum) # insert prefix-sum
ps.insert(bisect.bisect_left(ps,psum),psum)
#
print(best)
```
這題Python要做到$O(n\log n)$的時間複雜度,可以用一個資料結構SoredList,但它不在標準函式庫內,我們就不多做介紹,有興趣的讀者自行上網查詢用法,以下是範例程式。
```python=
# p_2_11, using SortedList, O(nlogn)
from sortedcontainers import SortedList
n, k = [int(x) for x in input().split()]
ps = SortedList([0]) # previous prefix_sum
best = -1000000000 # -oo
psum = 0 # current prefix sum
for v in (int(x) for x in input().split()):
psum += v # //prefix-sum(r)
i = ps.bisect_left(psum-k)
if i < len(ps): # found
best = max(best, psum-ps[i]) # currently best
ps.add(psum)
#
print(best)
```
以下的習題是二維版本的類似題。給一個提示,對每一個列的可能範圍$[i, j]$,把第$i$列到第$j$列”黏”起來變成一個一維陣列,然後用P-2-11的方法去做,這樣的時間複雜度會是$O(M^2 N\log(N))$。
---
##### Q-2-12. 最接近的子矩陣和 (108高中全國賽) (\*)
輸入一個整數二維矩陣$A[M][N]$,另外給了一個整數$K$,請計算哪一個子矩陣的和,也就是對所有$1\le i\le j\le M$ and $1\le p\le q\le N$, $\sum_{s=i}^{j}\sum_{t=p}^{q} A[s][t]$ 最接近$K$而不超過$K$。$M \le 10$且$M^2 \times N \le 100000$,每一個整數的絕對值不超過$3000$。
Time limit: 2 秒
輸入格式:每筆測資的第一行有ㄧ個正整數$K$;第二行有兩個正整數$M$ 與 $N$。接下來,由上而下,從左至右,有$M$行輸入,每一行有$N$個整數,每一個整數的絕對值不超過3,000,代表$A[s][t]$,同行整數間以空格隔開。
輸出格式:在所有子矩陣和中,最接近K但不超過K的和。
範例輸入:
6
2 3
-1 -1 1
2 2 2
範例輸出:
6
註:本題超出APCS範圍。
---
以下習題是108年高中全國賽某一題的核心要算的,是快速冪的一個應用。
---
##### Q-2-13. 無理數的快速冪 (108高中全國賽, simplifed)
若$s+t\sqrt{2}= (x+y\sqrt{2})^n$,其中$x, y, s, t$均為正整數,輸入$x$,$y$與$n$,請計算並輸出$s$與$t$除以$p$的餘數。$p=10^9+9$且$x,y,n < p$。
Time limit: 1秒
輸入格式:一行含三個正整數,依序為$x, y$與$n$,以空格隔開。
輸出格式:$s$與$t$,中間空一格。
範例輸入:
2 3 2
範例輸出:
22 12
---
下一題是108高中全國賽的題目,有多種解法,題目需要一些思考,但其實有個漂亮的解法是排序與遞迴的聯合運用,不需要特別的資料結構與方法。這一題有點難,如果做得出來應該已經超過APCS五級分的程度了。
---
##### Q-2-14. 水槽 (108高中全國賽) (@@)
我們用一排 $n$ 個擋板建造水槽。擋板的寬度為1,高度為正整數且均不相同,水槽前後是兩片長寬均為無限大的玻璃板 (見下圖例)。相鄰擋板的距離都是1,故相鄰二擋板之間會形成底面積1平方的水槽。
擋板由左而右依序由 0 到 $n – 1$編號,第 $i$ 及 $i + 1$ 檔板中間的水槽稱為水槽$i$。現在將總量為 $w$ 立方公分的水緩緩注入水槽$i$。注意水量可能溢出到別的水槽,但是由於所有擋板高度都不同,所以每當溢出時,只會先從一個方向溢出。請計算將總量為$w$立方公分的水緩緩注入水槽$i$後,所有水槽的水深。本題最左的擋板與最右的擋板是所有擋板中最高的兩個,並且保證欲注入的水不會溢出到左右邊界之外;另外,所有水槽的最後水深一定都是整數。以下圖為例,於水槽2注入17立方公分的水後,各水槽的水深依序為5, 5, 5, 1, 1, 0。

Time limit: 1秒
輸入格式:第一行有三個正整數$n\ (3 \le n \le 10^5)$、$i\ (0 \le i \le n – 2)$ 和 $w\ (1 \le w \le 10^{12})$,分別代表擋板數、注水水槽編號,及水量。第二行有 $n$ 個以空白間隔的正整數,代表由左到右擋板的高度,每個擋板高度為正整數且不超過$10^9$)。請注意注水量可能超過一個32-bit整數的範圍。。
輸出格式:輸出爲一行,共 $n – 1$ 個整數,依序代表各個水槽水深,數字之間以一個空白間隔。
範例輸入:
8 3 27
9 7 5 3 4 6 8 10
範例輸出:
0 6 6 6 6 3 0
註:本題Python如果寫遞迴版本可能會有超過遞迴深度與記憶體不足的問題。
---
Q_2_14解題提示:我們維護一個水槽高度尚未被決定的區間[left,right),區間以外的水槽高度都已經確定。
假設目前的水量為$water$以及注水點為$inp$。如果區間只剩一個水槽,該水槽高度即為水量$water$,結束。否則,在區間中找出最高的隔板,假設此隔板高度為$maxh$。區間被此隔板分成左右兩邊,考慮下面三種情形:
* $water >=(right-left)\times maxh$,水量足以讓兩邊都至少$maxh$,區間內所有水槽高度皆是相同的平均值,結束。
* 水量不足以越過輸入這一邊,那麼,另外一邊一定不會有水,縮小區間繼續計算。
* 水量會越過輸入這一邊,那麼,輸入這一邊的水槽高度一定都是$maxh$,把水量扣掉後,所小區間計算另外一邊。
為了要找出區間內的最高隔板,我們可以一開始把所有隔板的依照高度從高到低排列,每次要找某區間最高隔板時,我們檢視目前最高的隔板,如果在我們要的區間,那就找到了;如果不在我們要的區間,這個隔板以後也不會用到(想想為何),可以直接丟掉。因此,只要一開始把隔板排序後依序往後找就可以了,不需要特殊的資料結構,但是要把隔板的高度與位置一起存,所以需要一個兩個欄位資料的排序。
下一題是APCS考古題。
---
##### P-2-15. 圓環出口 (APCS202007)
有$n$ 個房間排列成一個圓環,以順時針方向由0到$n - 1$編號。玩家只能順時針方向依序通過這些房間。每當離開第$i$號房間進入下一個房間時,即可獲得$p(i)$點。玩家必須依序取得$m$把鑰匙,鑰匙編號由$0$至$m-1$,兌換編號$i$的鑰匙所需的點數為$Q(i)$。一旦玩家手中的點數達到$Q(i)$就會自動獲得編號$i$的鑰匙,而且手中所有的點數就會被「全數收回」,接著要再從當下所在的房間出發,重新收集點數兌換下一把鑰匙。遊戲開始時,玩家位於0號房。請計算玩家拿到最後一把鑰匙時所在的房間編號。
以下是一個例子。有7個房間,$p(i)$依序是$(2, 1, 5, 4, 3, 5, 3)$,其中0號房間的點數是2。假設所需要的鑰匙為3把,$Q(i)$ 依序是$(8, 9, 12)$。從0號房出發,在「離開2號房,進入3號房」時,獲得恰好2+1+5 = 8點,因此在進入3號房時玩家兌換到0號鑰匙;接著從3號房開始繼續累積點數,直到「離開5號房,進入6號房」時,手中的點數為12,於是在進入6號房時獲得1號鑰匙,手中點數再次被清空。最後,從6號房出發,直到「離開3號房,進入4號房」時,方可獲得至少12點的點數,來兌換最後一把鑰匙。因此,拿到最後一把鑰匙時所在的房間編號為4。
Time limit: 1秒
輸入格式:第一行有兩個正整數$n$ 與$m$,第二行有$n$個正整數,依序是各房間的點數$p(i)$,第三行有$ m$ 個正整數依序是各鑰匙需要的點數$Q(i)$,。同一行連續二數字間以空白隔開。$n$不超過2e5,$m$不超過2e4,$p(i)$總和不超過1e9,$Q(i)$不超過所有$p(i)$總和。
輸出格式:輸出拿到最後一把鑰匙時所在的房間編號。
範例輸入:
7 3
2 1 5 4 3 5 3
8 9 12
範例輸出:
4
---
簡單說,這個題目就是每次要在一個正整數陣列中,對於某個開始位置,找到一個最小的區間,滿足區間和大於等於某輸入$Q(i)$。關於圓環,有兩個簡單的處理方式,一個是判斷超過結尾時,從頭開始;另外一種方式是將陣列重複兩次,每次找到後除以n取餘數。最簡單的搜尋方法就是一個一個往下找,程式也很好寫,但是效率不夠好,以下是範例程式。
```python=
n, m = map(int, input().split())
a = [int(x) for x in input().split()]
q = [int(x) for x in input().split()]
room = 0 # current room
for qi in q: # each required qi
while qi > 0:
qi -= a[room]
room = (room+1)%n # next of n-1 is 0
#
print(room)
```
要有效率的搜尋,腦海中浮現的自然是二分搜,但是二分搜必須是在單調序列上才能使用,而且我們要找的是區間和。在P-1-6中我們介紹過前綴和(Prefix-sum),如果我們把所有房間的點數改成計算出前綴和,那麼正數的前綴和必然是遞增的,而我們要找的區間和也自然的轉換成找一個前綴和。關於二分搜,我們推薦前面介紹的一路往前跳的寫法,以下是範例程式。第4行的迴圈計算前綴和,然後第8行的迴圈處理每一次需求的點數,在第9行,如果目前房間不是0,就將所需的點數$qi$加上前一個房間的前綴和,這就是我們要搜尋的前綴和;因為本題是圓環,在第10行我們判斷是否需求會超過陣列尾端從0開始。接著進行二分搜,因為我們要始終保持$a[room]<qi$這個特性,跳躍結束時所在位置是最後一個小於$qi$的,所以根據題意,要加2才是下一個開始位置。要小心在二分搜開始往前跳之前,需要先檢查目前位置是否已經滿足需求。
```python=
n, m = map(int, input().split())
a = [int(x) for x in input().split()]
q = [int(x) for x in input().split()]
for i in range(1,n): # prefix sum
a[i] += a[i-1]
room = 0
total = a[n-1]
for qi in q:
if room>0: qi += a[room-1] # base
if qi>total: # exceed last, start from 0
room = 0
qi -= total
if a[room] >= qi: # current room
room = (room+1)%n
continue
# binary search the last <q[i]
jump = (n-room)//2
while jump>0:
while room+jump<n and a[room+jump]<qi:
room += jump
jump //= 2
#
room = (room + 2) %n
#end for
print(room)
```
我們當然也可以用內建二分搜bisect.bisect_left來做二分搜,以下是範例程式,寫法類似就不多做說明。
```python=
import bisect
n, m = map(int, input().split())
a = [int(x) for x in input().split()]
q = [int(x) for x in input().split()]
for i in range(1,n): # prefix sum
a[i] += a[i-1]
room = 0
total = a[n-1]
for qi in q:
if room>0: qi += a[room-1] # base
if qi>total: # exceed last, start from 0
room = 0
qi -= total
# search the first >= q[i]
if a[room] < qi:
room = bisect.bisect_left(a,qi)
room = (room+1)%n
#end for
print(room)
```
**End of Chapter 2**