チンチラのフンみたいなもの

毎日の学んだことを書いていきます

競プロ典型90問 010 - Score Sum Queries(★2)

問題はコチラから

問題文

ABC 大学には N 人の一年生が在籍しています。クラスは 2 つあり、学籍番号 i 番の生徒のクラスは Ci組です。今日は期末試験が返却され、学籍番号 i 番の生徒の点数は Pi点でした。
以下の形式の質問が Q 個与えられます。j=1,2,…,Q それぞれについて答えてください。
・学籍番号 Lj∼Rj番の 1 組生徒における、期末試験点数の合計
・学籍番号 Lj∼Rj番の 2 組生徒における、期末試験点数の合計
・これら 2 つの値をそれぞれ求めよ。

制約

  • 1≤N≤100000
  • 1≤Ci ≤2
  • 0≤Pi≤100
  • 1≤Q≤100000
  • 1≤Lj≤Rj≤N
  • 入力は全て整数

私のコード

 単純にたすだけじゃんと思って書いたのだが、TLEになる。Time Limit Exceededの略で実行時間超過を意味する。今回だと2秒以上かかってしまったためACにならなかった。この問題は2重ループをするとよく起こるみたい。実際私もfor文の中にfor文を入れるということをしている。
 ならfor文を一つ減らそうと思った。一番外側のfor文は答えをprintする部分だから解消する意味は薄そうに思った。ならもう一つの方を変えようと思ったが、なかなか良い方法が思いつかない。そこで今回は他の人のコードを見て勉強させて頂く!

let a = Int(readLine()!)!
var data = [[Int]]()
for _ in 1 ... a {data.append(readLine()!.split(separator: " ").map{Int($0)!})}
let q = Int(readLine()!)!
var question = [[Int]]()
for _ in 1 ... q {question.append(readLine()!.split(separator: " ").map{Int($0)!})}
for j in 0 ..< q {
    var sum1 = 0
    var sum2 = 0
    for i in question[j][0]-1 ..< question[j][1] {
        if data[i][0] == 1{
            sum1 += data[i][1]
        }else{
            sum2 += data[i][1]
        }
    }
    print("\(sum1) \(sum2)")
}

他の人のコード

func readInt() -> Int {
    Int(readLine()!)!
}

func readIntArray() -> [Int] {
    readLine()!.split(separator: " ").map { Int(String($0))! }
}

let n = readInt()
var class1Sum = [Int](repeating: 0, count: n + 1)
var class2Sum = [Int](repeating: 0, count: n + 1)

for i in 1...n {
    let cp = readIntArray()
    let c = cp[0]
    let p = cp[1]
    if c == 1 {
        class1Sum[i] = class1Sum[i - 1] + p
        class2Sum[i] = class2Sum[i - 1]
    } else if c == 2 {
        class1Sum[i] = class1Sum[i - 1]
        class2Sum[i] = class2Sum[i - 1] + p
    }
}

let q = readInt()
for _ in 0..<q {
    let lr = readIntArray()
    let l = lr[0]
    let r = lr[1]
    print("\(class1Sum[r] - class1Sum[l - 1]) \(class2Sum[r] - class2Sum[l - 1])")
}

改善点

  • [ [ Int ] ]型の変数を宣言しなくても、for文の中で[Int]型の変数を宣言することで処理数を減らせる
  • [Int] (repeating: 0, count: n + 1)を使うことで、var sum1 = 0、var sum2 = 0といった宣言をしなくて済む

解説