累積和の訓練

qiita.com

こちらの記事ベースで学習。

AOJ 0516 - 最大の和 (JOI 2006 本選 A)

judge.u-aizu.ac.jp

この問題、入力をどう書くべきかいまいち分からなかったのでテストコードが通れば良しとする。

テストコード

import unittest
import main


class MainTest(unittest.TestCase):
    def setUp(self):
        pass

    def tearDown(self):
        pass

    def test_1(self):
        self.assertEqual(11, main.main(5, 3, [2, 5, -4, 10, 3]))


if __name__ == "__main__":
    unittest.main()

提出（してないけど）コード

def main(N, K, a_list):
    s_list = [0]
    for i, _ in enumerate(a_list):
        s_list.append(s_list[i] + a_list[i])

    res = float('-inf')
    for i in range(N-K):
        val = s_list[K+i] - s_list[i]
        if val > res:
            res = val

    return res

abc122_c Get AC

atcoder.jp

テストコード

import unittest
import main


class MainTest(unittest.TestCase):
    def setUp(self):
        pass

    def tearDown(self):
        pass

    def test_1(self):
        self.assertEqual([2, 0, 3], main.main(
            8, 3, "ACACTACG", [(3, 7), (2, 3), (1, 8)]))


if __name__ == "__main__":
    unittest.main()

提出コード (TLE)

愚直に全探索。

from math import ceil


def main(N, Q, S, lr_list):
    res = []
    for lr in lr_list:
        l, r = lr
        res.append(S[l-1:r].count("AC"))
    return res


if __name__ == "__main__":
    N, Q = list(map(int, input().split()))
    S = input()
    lr_list = []
    for _ in range(Q):
        lr_list.append(tuple(map(int, input().split())))
    res = main(N, Q, S, lr_list)
    for r in res:
        print(r)

当然TLE。

提出コード (TLE)

累積和を取り入れる。

def main(N, Q, S, lr_list):
    s_list = [0]
    for i in range(N):
        if i+1 < N and list(S)[i] == 'A' and list(S)[i+1] == 'C':
            s_list.append(s_list[i] + 1)
        else:
            s_list.append(s_list[i])

    res_list = []
    for lr in lr_list:
        l, r = lr
        l -= 1
        r -= 1
        res_list.append(s_list[r] - s_list[l])

    return res_list


if __name__ == "__main__":
    N, Q = list(map(int, input().split()))
    S = input()
    lr_list = []
    for _ in range(Q):
        lr_list.append(tuple(map(int, input().split())))
    res = main(N, Q, S, lr_list)
    for r in res:
        print(r)

しかしまだTLE。 forを回す都度 list(S) を行うのがボトルネックになっていたようだ。

提出コード (AC)

def main(N, Q, S, lr_list):
    chars = list(S) // ここで文字列を配列にしておく

    s_list = [0]
    for i in range(N):
        if i+1 < N and chars[i] == 'A' and chars[i+1] == 'C':
            s_list.append(s_list[i] + 1)
        else:
            s_list.append(s_list[i])

    res_list = []
    for lr in lr_list:
        l, r = lr
        l -= 1
        r -= 1
        res_list.append(s_list[r] - s_list[l])

    return res_list


if __name__ == "__main__":
    N, Q = list(map(int, input().split()))
    S = input()
    lr_list = []
    for _ in range(Q):
        lr_list.append(tuple(map(int, input().split())))
    res = main(N, Q, S, lr_list)
    for r in res:
        print(r)

これでAC。